Spaces:
Running on Zero
Running on Zero
| """Fit a whole photo into the model's square input instead of cropping it. | |
| The model takes 448x448. `letterbox` scales the photo so its LONG side is 448 | |
| and fills the rest of the square by mirroring the photo's own edge (the model | |
| never saw flat-coloured bars, and the test-time rotation already feeds it | |
| reflected borders). After the model runs, `crop_to_box` cuts the padding off | |
| the maps again, so nothing of the photo is lost and nothing invented is shown. | |
| """ | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from PIL import Image | |
| def letterbox(img: Image.Image, size: int): | |
| """Return (square PIL image, (left, top, w, h) of the real photo in it).""" | |
| img = img.convert("RGB") | |
| w, h = img.size | |
| scale = size / float(max(w, h)) | |
| nw, nh = max(1, int(round(w * scale))), max(1, int(round(h * scale))) | |
| arr = np.asarray(img.resize((nw, nh), Image.BICUBIC)) | |
| left, top = (size - nw) // 2, (size - nh) // 2 | |
| pad = ((top, size - nh - top), (left, size - nw - left), (0, 0)) | |
| # 'symmetric' (edge pixel included) is safe for any pad width, even when | |
| # the padding is longer than the photo itself (a very elongated panorama). | |
| out = np.pad(arr, pad, mode="symmetric") | |
| return Image.fromarray(out), (left, top, nw, nh) | |
| def crop_to_box(maps: torch.Tensor, box) -> torch.Tensor: | |
| """Cut [B, C, S, S] maps down to the photo's own pixels.""" | |
| left, top, w, h = box | |
| return maps[:, :, top:top + h, left:left + w] | |
| def resize_maps(maps: torch.Tensor, width: int, height: int) -> torch.Tensor: | |
| """Resize 7-channel maps for display and keep the normals unit length.""" | |
| out = F.interpolate(maps.float(), size=(height, width), mode="bicubic", | |
| align_corners=False).clamp(-1.0, 1.0) | |
| n = out[:, 4:7] | |
| out[:, 4:7] = n / n.norm(dim=1, keepdim=True).clamp_min(1e-6) | |
| return out | |
| def display_size(w: int, h: int, longest: int = 1024): | |
| """Output size: the photo's own shape, long side at most `longest`.""" | |
| s = min(1.0, longest / float(max(w, h))) | |
| return max(1, int(round(w * s))), max(1, int(round(h * s))) | |