Spaces:
Running on Zero
Running on Zero
| """In-plane rotation of photos and maps, for test-time rotation averaging. | |
| Conventions (match the training augmentation, `MarigoldV2Dataset.__getitem__`): | |
| * a positive angle rotates the picture counter-clockwise (like PIL's rotate); | |
| * the maps are `[0:3] albedo | [3] depth | [4:7] normal`, with the normal in the | |
| camera frame x right, y UP, z toward the viewer. Albedo and depth are | |
| invariant under a roll about the optical axis, so only their positions move; | |
| the normal's (nx, ny) rotate with the picture by R(+theta) = | |
| [[c, -s], [s, c]]. | |
| `rotate_image` resamples with a reflected border (no empty corners, which the | |
| model never saw); `unrotate_maps` returns the maps back in the original frame | |
| together with the mask of pixels the rotated canvas actually covered. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| import torch | |
| import torch.nn.functional as F | |
| def _theta(angle_deg: float, device, dtype) -> torch.Tensor: | |
| """affine_grid matrix [1, 2, 3]: output pixel -> input sample position, so | |
| that the picture content turns counter-clockwise by `angle_deg` (image | |
| coordinates: x right, y down).""" | |
| a = math.radians(angle_deg) | |
| c, s = math.cos(a), math.sin(a) | |
| return torch.tensor([[[c, -s, 0.0], [s, c, 0.0]]], device=device, | |
| dtype=dtype) | |
| def rotate_image(x: torch.Tensor, angle_deg: float, mode: str = "bicubic", | |
| padding: str = "reflection") -> torch.Tensor: | |
| """Rotate [B, C, H, W] counter-clockwise by `angle_deg` about the centre. | |
| Sampled in float32 with autocast off: the sampling grid in bf16 has only ~3 | |
| significant digits, which at 448 px is a shift of a pixel or more (found | |
| when the same code under `torch.autocast` gave outputs 0.5 off).""" | |
| b = x.shape[0] | |
| with torch.autocast(device_type=x.device.type, enabled=False): | |
| xf = x.float() | |
| theta = _theta(angle_deg, x.device, torch.float32).expand(b, -1, -1) | |
| grid = F.affine_grid(theta, list(xf.shape), align_corners=False) | |
| out = F.grid_sample(xf, grid, mode=mode, padding_mode=padding, | |
| align_corners=False) | |
| return out.to(x.dtype) | |
| def _reorient_normals(maps: torch.Tensor, angle_deg: float) -> torch.Tensor: | |
| """Apply R(angle) to the normal's (nx, ny); nz is unchanged.""" | |
| a = math.radians(angle_deg) | |
| c, s = math.cos(a), math.sin(a) | |
| out = maps.clone() | |
| nx, ny = maps[:, 4], maps[:, 5] | |
| out[:, 4] = c * nx - s * ny | |
| out[:, 5] = s * nx + c * ny | |
| return out | |
| def rotate_maps(maps: torch.Tensor, angle_deg: float): | |
| """Rotate 7ch maps counter-clockwise and reorient the normals. | |
| Returns (maps, valid) where `valid` [B, 1, H, W] marks canvas pixels that | |
| came from inside the original frame (the rest are zero).""" | |
| ones = torch.ones_like(maps[:, :1]) | |
| r = rotate_image(maps, angle_deg, "bilinear", "zeros") | |
| valid = rotate_image(ones, angle_deg, "bilinear", "zeros") > 0.999 | |
| return _reorient_normals(r, angle_deg), valid | |
| def unrotate_maps(maps: torch.Tensor, angle_deg: float): | |
| """Maps predicted on a photo rotated by `angle_deg`, brought back to the | |
| original frame (rotate by -angle, normals reoriented by R(-angle)). | |
| Returns (maps, valid): `valid` marks original-frame pixels the rotated | |
| canvas covered (the corners are lost for rotations of a few degrees).""" | |
| return rotate_maps(maps, -angle_deg) | |