Spaces:
Running on Zero
Running on Zero
File size: 2,107 Bytes
b3cdd04 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | """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)))
|