(x)
| 94 | |
| 95 | |
| 96 | def grad(x): |
| 97 | # x.shape : n, c, h, w |
| 98 | diff_x = x[..., 1:, 1:] - x[..., 1:, :-1] |
| 99 | diff_y = x[..., 1:, 1:] - x[..., :-1, 1:] |
| 100 | mag = diff_x**2 + diff_y**2 |
| 101 | # angle_ratio |
| 102 | angle = torch.atan(diff_y / (diff_x + 1e-10)) |
| 103 | return mag, angle |
| 104 | |
| 105 | |
| 106 | def grad_mask(mask): |