| """Aspect-preserving colourisation with a compact semantic model. |
| |
| Input and output are PIL images. No teacher, critic or second learned model is |
| loaded. Original lightness and alpha are retained; out-of-gamut chroma is reduced. |
| """ |
| import json |
| import math |
| from pathlib import Path |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from PIL import Image, ImageOps |
| from skimage.color import rgb2lab |
| from semantic_model import load_semantic |
|
|
| MAX_PIXELS = 12_000_000 |
|
|
| def load_colorizer(path, device='cpu'): |
| model = load_semantic(path, device) |
| count = sum(p.numel() for p in model.parameters()) |
| if count >= 4_000_000: |
| raise ValueError('Model exceeds the four-million-parameter limit') |
| return model |
|
|
| def box_mean(x, r): |
| return F.avg_pool2d(x, 2*r+1, 1, r, count_include_pad=False) |
|
|
| def prepare(image, size=256): |
| if not isinstance(image, Image.Image): |
| raise TypeError('Expected a PIL image') |
| if not 128 <= int(size) <= 512: |
| raise ValueError('Input size must be between 128 and 512') |
| image = ImageOps.exif_transpose(image) |
| if image.width * image.height > MAX_PIXELS: |
| raise ValueError('Please resize the image to at most 12 megapixels') |
| alpha = np.asarray(image.getchannel('A')).copy() if 'A' in image.getbands() else None |
| rgb = np.asarray(image.convert('RGB'), dtype=np.float32) / 255 |
| light = rgb2lab(rgb)[..., 0].astype(np.float32) |
| scale = min(int(size)/max(image.size), 1.) |
| shape = (max(8, round(image.height*scale)), max(8, round(image.width*scale))) |
| x = torch.from_numpy(light)[None,None]/50-1 |
| small = F.interpolate(x, size=shape, mode='bilinear', align_corners=False, antialias=True) |
| return light, alpha, small |
|
|
| @torch.inference_mode() |
| def chroma_coefficients(model, small, radius=8): |
| if radius not in (0,4,8,12,16): |
| raise ValueError('Unsupported smoothing radius') |
| device = next(model.parameters()).device |
| L = small.to(device) |
| ab = model(L).float() |
| if not torch.isfinite(ab).all(): |
| raise RuntimeError('Model returned non-finite colours') |
| if radius == 0: |
| return torch.zeros_like(ab).cpu(), ab.cpu() |
| guide = (L.float()+1)/2 |
| mi, mp = box_mean(guide,radius), box_mean(ab,radius) |
| var = (box_mean(guide*guide,radius)-mi*mi).clamp_min(0) |
| cov = box_mean(guide*ab,radius)-mi*mp |
| a = cov/(var+.001) |
| b = mp-a*mi |
| |
| return box_mean(a,radius).cpu(), box_mean(b,radius).cpu() |
|
|
| def _linear_rgb(light, ab): |
| fy = (light+16)/116 |
| f = np.stack([fy+ab[...,0]/500,fy,fy-ab[...,1]/200],axis=-1) |
| xyz = np.where(f>6/29,f**3,(f-4/29)*(3*(6/29)**2)) |
| xyz *= np.array([.95047,1.,1.08883],np.float32) |
| matrix = np.array([[3.24048134,-1.53715152,-.49853633],[-.96925495,1.87599,.04155593],[.05564664,-.20404134,1.05731107]],np.float32) |
| return xyz @ matrix.T |
|
|
| def render(light, alpha, coefficients, saturation=1.): |
| saturation = float(saturation) |
| if not math.isfinite(saturation) or not 0 <= saturation <= 1.5: |
| raise ValueError('Colour strength must be between 0 and 1.5') |
| a,b = [F.interpolate(v.float(),size=light.shape,mode='bilinear',align_corners=False)[0].permute(1,2,0).numpy() for v in coefficients] |
| ab = (a*(light[...,None]/100)+b)*saturation |
| |
| linear = _linear_rgb(light,ab) |
| invalid = ((linear < -1e-5)|(linear > 1+1e-5)).any(-1) |
| if invalid.any(): |
| L = light[invalid]; colors=ab[invalid];lo=np.zeros(len(L),np.float32);hi=np.ones(len(L),np.float32) |
| for _ in range(9): |
| mid=(lo+hi)/2; candidate=_linear_rgb(L,colors*mid[:,None]) |
| valid=((candidate>=-1e-5)&(candidate<=1+1e-5)).all(-1) |
| lo=np.where(valid,mid,lo);hi=np.where(valid,hi,mid) |
| ab[invalid]=colors*lo[:,None] |
| linear[invalid]=_linear_rgb(L,ab[invalid]) |
| linear=np.clip(linear,0,1) |
| rgb=np.where(linear<=.0031308,12.92*linear,1.055*np.power(linear,1/2.4)-.055) |
| pixels=np.uint8(np.clip(np.rint(rgb*255),0,255)) |
| if alpha is not None: pixels=np.concatenate([pixels,alpha[...,None]],axis=-1) |
| return Image.fromarray(pixels) |
|
|
| def colorize(model, image, size=256, radius=8, saturation=1.): |
| light,alpha,small=prepare(image,size) |
| return render(light,alpha,chroma_coefficients(model,small,radius),saturation) |
|
|
| def main(): |
| import argparse |
| parser=argparse.ArgumentParser(description='Compact photo colouriser') |
| parser.add_argument('input');parser.add_argument('output') |
| parser.add_argument('--model',default='.');parser.add_argument('--device',default='cpu') |
| parser.add_argument('--size',type=int,default=256);parser.add_argument('--saturation',type=float,default=1.) |
| args=parser.parse_args() |
| model=load_colorizer(args.model,args.device) |
| with Image.open(args.input) as image: |
| colorize(model,image,args.size,saturation=args.saturation).save(args.output) |
|
|
| if __name__=='__main__':main() |
|
|
