File size: 2,975 Bytes
704fa80 | 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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """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()
|