| 1074 | self.kernel_diff, self.kernel_smooth = self._get_kernel(kernel_size, dtype) |
| 1075 | |
| 1076 | def _get_kernel(self, size, dtype) -> tuple[torch.Tensor, torch.Tensor]: |
| 1077 | if size < 3: |
| 1078 | raise ValueError(f"Sobel kernel size should be at least three. {size} was given.") |
| 1079 | if size % 2 == 0: |
| 1080 | raise ValueError(f"Sobel kernel size should be an odd number. {size} was given.") |
| 1081 | |
| 1082 | kernel_diff = torch.tensor([[[-1, 0, 1]]], dtype=dtype) |
| 1083 | kernel_smooth = torch.tensor([[[1, 2, 1]]], dtype=dtype) |
| 1084 | kernel_expansion = torch.tensor([[[1, 2, 1]]], dtype=dtype) |
| 1085 | |
| 1086 | if self.normalize_kernels: |
| 1087 | if not dtype.is_floating_point: |
| 1088 | raise ValueError( |
| 1089 | f"`dtype` for Sobel kernel should be floating point when `normalize_kernel==True`. {dtype} was given." |
| 1090 | ) |
| 1091 | kernel_diff /= 2.0 |
| 1092 | kernel_smooth /= 4.0 |
| 1093 | kernel_expansion /= 4.0 |
| 1094 | |
| 1095 | # Expand the kernel to larger size than 3 |
| 1096 | expand = (size - 3) // 2 |
| 1097 | for _ in range(expand): |
| 1098 | kernel_diff = F.conv1d(kernel_diff, kernel_expansion, padding=2) |
| 1099 | kernel_smooth = F.conv1d(kernel_smooth, kernel_expansion, padding=2) |
| 1100 | |
| 1101 | return kernel_diff.squeeze(), kernel_smooth.squeeze() |
| 1102 | |
| 1103 | def __call__(self, image: NdarrayOrTensor) -> torch.Tensor: |
| 1104 | image_tensor = convert_to_tensor(image, track_meta=get_track_meta()) |