def forward(self, prediction, target, preview = False): # helper to remove potential nan in labels def nan_helper(y): return torch.isnan(y), lambda z: z.nonzero()[0] loss_value = 0 loss_value_2 = 0 diff = prediction - target _,_,H,W = target.shape upsample = torch.nn.Upsample(size=(2*H,2*W), mode='bicubic', align_corners=True) record = [] for m in self.multi_scales: # input and type are of the type [B x C x H x W] if preview: record.append(upsample(sobel(m(diff)))) else: # Use kornia spatial gradient computation delta_diff = spatial_gradient(m(diff)) is_nan = torch.isnan(delta_diff) is_not_nan_sum = (~is_nan).sum() # output of kornia spatial gradient is [B x C x 2 x H x W] loss_value += torch.abs(delta_diff[~is_nan]).sum()/is_not_nan_sum*target.shape[0]*2 # * batch size * 2 (because kornia spatial product has two outputs). # replaces the following line to be able to deal with nan's. # loss_value += torch.abs(delta_diff).mean(dim=(3,4)).sum() if preview: return record else: return (loss_value/self.num_scales)
def forward(self, input: torch.Tensor, axis=None) -> torch.Tensor: # type: ignore if axis is None: axis = self.axis if not torch.is_tensor(input): raise TypeError("Input type is not a torch.Tensor. Got {}".format( type(input))) if not len(input.shape) == 4: raise ValueError( "Invalid input shape, we expect BxCxHxW. Got: {}".format( input.shape)) # comput the x/y gradients edges: torch.Tensor = spatial_gradient(input, normalized=self.normalized) # unpack the edges gx: torch.Tensor = edges[:, :, 0] gy: torch.Tensor = edges[:, :, 1] # compute gradient maginitude magnitude: torch.Tensor = torch.sqrt(gx * gx + gy * gy) if axis is None: return magnitude elif axis in {3, -1}: return gx elif axis in {2, -2}: return gy
def derivat(img, mode='sobel'): if mode == 'scharr': # https://en.wikipedia.org/wiki/Sobel_operator#Alternative_operators k_scharr = torch.Tensor([[[-0.183,0.,0.183], [-0.634,0.,0.634], [-0.183,0.,0.183]], [[-0.183,-0.634,-0.183], [0.,0.,0.], [0.183,0.634,0.183]]]) k_scharr = k_scharr.unsqueeze(1).tile((1,3,1,1)).cuda() return 0.2 * torch.mean(torch.abs(F.conv2d(img, k_scharr))) elif mode == 'sobel': # https://kornia.readthedocs.io/en/latest/filters.html#edge-detection return torch.mean(torch.abs(spatial_gradient(img))) else: # trivial hack dx = torch.mean(torch.abs(img[:,:,:,1:] - img[:,:,:,:-1])) dy = torch.mean(torch.abs(img[:,:,1:,:] - img[:,:,:-1,:])) return 0.5 * (dx+dy)
def depth_to_normals(depth: torch.Tensor, camera_matrix: torch.Tensor, normalize_points: bool = False) -> torch.Tensor: """Compute the normal surface per pixel. Args: depth: image tensor containing a depth value per pixel with shape :math:`(B, 1, H, W)`. camera_matrix: tensor containing the camera intrinsics with shape :math:`(B, 3, 3)`. normalize_points: whether to normalise the pointcloud. This must be set to `True` when the depth is represented as the Euclidean ray length from the camera position. Return: tensor with a normal surface vector per pixel of the same resolution as the input :math:`(B, 3, H, W)`. Example: >>> depth = torch.rand(1, 1, 4, 4) >>> K = torch.eye(3)[None] >>> depth_to_normals(depth, K).shape torch.Size([1, 3, 4, 4]) """ if not isinstance(depth, torch.Tensor): raise TypeError( f"Input depht type is not a torch.Tensor. Got {type(depth)}.") if not (len(depth.shape) == 4 and depth.shape[-3] == 1): raise ValueError( f"Input depth musth have a shape (B, 1, H, W). Got: {depth.shape}") if not isinstance(camera_matrix, torch.Tensor): raise TypeError(f"Input camera_matrix type is not a torch.Tensor. " f"Got {type(camera_matrix)}.") if not (len(camera_matrix.shape) == 3 and camera_matrix.shape[-2:] == (3, 3)): raise ValueError(f"Input camera_matrix must have a shape (B, 3, 3). " f"Got: {camera_matrix.shape}.") # compute the 3d points from depth xyz: torch.Tensor = depth_to_3d(depth, camera_matrix, normalize_points) # Bx3xHxW # compute the pointcloud spatial gradients gradients: torch.Tensor = spatial_gradient(xyz) # Bx3x2xHxW # compute normals a, b = gradients[:, :, 0], gradients[:, :, 1] # Bx3xHxW normals: torch.Tensor = torch.cross(a, b, dim=1) # Bx3xHxW return F.normalize(normals, dim=1, p=2)
def canny( input: torch.Tensor, low_threshold: float = 0.1, high_threshold: float = 0.2, kernel_size: Tuple[int, int] = (5, 5), sigma: Tuple[float, float] = (1, 1), hysteresis: bool = True, eps: float = 1e-6, ) -> Tuple[torch.Tensor, torch.Tensor]: r"""Finds edges of the input image and filters them using the Canny algorithm. .. image:: _static/img/canny.png Args: input: input image tensor with shape :math:`(B,C,H,W)`. low_threshold: lower threshold for the hysteresis procedure. high_threshold: upper threshold for the hysteresis procedure. kernel_size: the size of the kernel for the gaussian blur. sigma: the standard deviation of the kernel for the gaussian blur. hysteresis: if True, applies the hysteresis edge tracking. Otherwise, the edges are divided between weak (0.5) and strong (1) edges. eps: regularization number to avoid NaN during backprop. Returns: - the canny edge magnitudes map, shape of :math:`(B,1,H,W)`. - the canny edge detection filtered by thresholds and hysteresis, shape of :math:`(B,1,H,W)`. Example: >>> input = torch.rand(5, 3, 4, 4) >>> magnitude, edges = canny(input) # 5x3x4x4 >>> magnitude.shape torch.Size([5, 1, 4, 4]) >>> edges.shape torch.Size([5, 1, 4, 4]) """ if not isinstance(input, torch.Tensor): raise TypeError("Input type is not a torch.Tensor. Got {}".format(type(input))) if not len(input.shape) == 4: raise ValueError("Invalid input shape, we expect BxCxHxW. Got: {}".format(input.shape)) if low_threshold > high_threshold: raise ValueError( "Invalid input thresholds. low_threshold should be smaller than the high_threshold. Got: {}>{}".format( low_threshold, high_threshold ) ) if low_threshold < 0 and low_threshold > 1: raise ValueError( "Invalid input threshold. low_threshold should be in range (0,1). Got: {}".format(low_threshold) ) if high_threshold < 0 and high_threshold > 1: raise ValueError( "Invalid input threshold. high_threshold should be in range (0,1). Got: {}".format(high_threshold) ) device: torch.device = input.device dtype: torch.dtype = input.dtype # To Grayscale if input.shape[1] == 3: input = rgb_to_grayscale(input) # Gaussian filter blurred: torch.Tensor = gaussian_blur2d(input, kernel_size, sigma) # Compute the gradients gradients: torch.Tensor = spatial_gradient(blurred, normalized=False) # Unpack the edges gx: torch.Tensor = gradients[:, :, 0] gy: torch.Tensor = gradients[:, :, 1] # Compute gradient magnitude and angle magnitude: torch.Tensor = torch.sqrt(gx * gx + gy * gy + eps) angle: torch.Tensor = torch.atan2(gy, gx) # Radians to Degrees angle = rad2deg(angle) # Round angle to the nearest 45 degree angle = torch.round(angle / 45) * 45 # Non-maximal suppression nms_kernels: torch.Tensor = get_canny_nms_kernel(device, dtype) nms_magnitude: torch.Tensor = F.conv2d(magnitude, nms_kernels, padding=nms_kernels.shape[-1] // 2) # Get the indices for both directions positive_idx: torch.Tensor = (angle / 45) % 8 positive_idx = positive_idx.long() negative_idx: torch.Tensor = ((angle / 45) + 4) % 8 negative_idx = negative_idx.long() # Apply the non-maximum suppresion to the different directions channel_select_filtered_positive: torch.Tensor = torch.gather(nms_magnitude, 1, positive_idx) channel_select_filtered_negative: torch.Tensor = torch.gather(nms_magnitude, 1, negative_idx) channel_select_filtered: torch.Tensor = torch.stack( [channel_select_filtered_positive, channel_select_filtered_negative], 1 ) is_max: torch.Tensor = channel_select_filtered.min(dim=1)[0] > 0.0 magnitude = magnitude * is_max # Threshold edges: torch.Tensor = F.threshold(magnitude, low_threshold, 0.0) low: torch.Tensor = magnitude > low_threshold high: torch.Tensor = magnitude > high_threshold edges = low * 0.5 + high * 0.5 edges = edges.to(dtype) # Hysteresis if hysteresis: edges_old: torch.Tensor = -torch.ones(edges.shape, device=edges.device, dtype=dtype) hysteresis_kernels: torch.Tensor = get_hysteresis_kernel(device, dtype) while ((edges_old - edges).abs() != 0).any(): weak: torch.Tensor = (edges == 0.5).float() strong: torch.Tensor = (edges == 1).float() hysteresis_magnitude: torch.Tensor = F.conv2d( edges, hysteresis_kernels, padding=hysteresis_kernels.shape[-1] // 2 ) hysteresis_magnitude = (hysteresis_magnitude == 1).any(1, keepdim=True).to(dtype) hysteresis_magnitude = hysteresis_magnitude * weak + strong edges_old = edges.clone() edges = hysteresis_magnitude + (hysteresis_magnitude == 0) * weak * 0.5 edges = hysteresis_magnitude return magnitude, edges