| """Minimal ONNX-only image example; fixed square network input. |
| |
| For aspect-preserving guided upsampling and gamut compression use inference.py. |
| """ |
| import argparse |
| import numpy as np |
| import onnxruntime as ort |
| from PIL import Image,ImageOps |
| from skimage.color import rgb2lab,lab2rgb |
|
|
| def colorize(image,model='colorizer.onnx'): |
| image=ImageOps.exif_transpose(image).convert('RGB') |
| if image.width*image.height>12_000_000:raise ValueError('Maximum 12 megapixels') |
| light=rgb2lab(np.asarray(image,np.float32)/255)[...,0].astype(np.float32) |
| resized=np.asarray(Image.fromarray(light).resize((256,256),Image.Resampling.BILINEAR)) |
| session=ort.InferenceSession(model,providers=['CPUExecutionProvider']) |
| ab=session.run(['ab'],{'L':(resized[None,None]/50-1).astype(np.float32)})[0][0] |
| ab=np.stack([np.asarray(Image.fromarray(channel).resize(image.size,Image.Resampling.BILINEAR)) for channel in ab],-1) |
| rgb=np.clip(lab2rgb(np.concatenate([light[...,None],ab],-1)),0,1) |
| return Image.fromarray(np.uint8(np.rint(rgb*255))) |
|
|
| if __name__=='__main__': |
| p=argparse.ArgumentParser();p.add_argument('input');p.add_argument('output');p.add_argument('--model',default='colorizer.onnx');a=p.parse_args() |
| colorize(Image.open(a.input),a.model).save(a.output) |
|
|
