| """Aspect-preserving inference; infer chroma globally and retain original luminance.""" |
| import argparse |
| import math |
| from pathlib import Path |
| import numpy as np |
| from PIL import Image, ImageOps |
| import torch |
| import torch.nn.functional as F |
| from skimage.color import rgb2lab, lab2rgb |
| from model import load_model |
| from spatial import guided_chroma |
|
|
| @torch.inference_mode() |
| def colorize(model, image, size=256, temperature=0.38, saturation=1.0, flip_tta=False, |
| guided_radius=8, guided_epsilon=.001): |
| if size < 8 or not math.isfinite(saturation) or saturation < 0: |
| raise ValueError("size must be >=8 and saturation finite and nonnegative") |
| image = ImageOps.exif_transpose(image).convert("RGB") |
| rgb = np.asarray(image, dtype=np.float32) / 255.0 |
| luminance = rgb2lab(rgb)[..., 0].astype(np.float32) |
| h, w = luminance.shape |
| scale = min(size / max(h, w), 1.0) |
| target = (max(8, round(h*scale)), max(8, round(w*scale))) |
| device = next(model.parameters()).device |
| L = torch.from_numpy(luminance)[None, None].to(device) / 50 - 1 |
| small = F.interpolate(L, size=target, mode="bilinear", align_corners=False, antialias=True) |
| logits = model(small) |
| if flip_tta: |
| logits = (logits + model(small.flip(-1)).flip(-1)) * 0.5 |
| ab = model.decode(logits, temperature) |
| ab = guided_chroma(small, ab, guided_radius, guided_epsilon) |
| ab = F.interpolate(ab, size=(h,w), mode="bilinear", align_corners=False) |
| ab = ab[0].permute(1,2,0).cpu().numpy() * saturation |
| lab = np.concatenate([luminance[...,None], ab], axis=-1) |
| result = np.clip(lab2rgb(lab), 0, 1) |
| return Image.fromarray(np.rint(result * 255).astype(np.uint8)) |
|
|
| def main(): |
| p = argparse.ArgumentParser(__doc__) |
| p.add_argument("images", nargs="+") |
| p.add_argument("--model", required=True) |
| p.add_argument("--revision", default=None) |
| p.add_argument("--size", type=int, default=256) |
| p.add_argument("--temperature", type=float, default=.38) |
| p.add_argument("--saturation", type=float, default=1) |
| p.add_argument("--flip-tta", action="store_true") |
| p.add_argument("--guided-radius",type=int,default=8,help="Chroma smoothing radius at model resolution;0 disables") |
| p.add_argument("--guided-epsilon",type=float,default=.001) |
| p.add_argument("--output-dir", default="colorized") |
| p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| a = p.parse_args() |
| torch.set_num_threads(min(torch.get_num_threads(),4)) |
| model = load_model(a.model, a.revision, a.device) |
| out = Path(a.output_dir); out.mkdir(parents=True, exist_ok=True) |
| for file in a.images: |
| with Image.open(file) as im: |
| result = colorize(model, im, a.size, a.temperature, a.saturation, a.flip_tta, |
| a.guided_radius,a.guided_epsilon) |
| dest = out / (Path(file).stem + "_colorized.png") |
| result.save(dest); print(dest) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
