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