| """Run the exported chroma pipeline without importing PyTorch.""" |
| import argparse |
| from pathlib import Path |
| import numpy as np |
| import onnxruntime as ort |
| from PIL import Image,ImageOps |
| from skimage.color import rgb2lab,lab2rgb |
|
|
| def load_session(path,threads=2): |
| opts=ort.SessionOptions();opts.intra_op_num_threads=threads;opts.inter_op_num_threads=1 |
| return ort.InferenceSession(str(path),sess_options=opts,providers=['CPUExecutionProvider']) |
|
|
| def colorize(session,image,size=256): |
| if size<8:raise ValueError('size must be at least8') |
| image=ImageOps.exif_transpose(image).convert('RGB') |
| L=rgb2lab(np.asarray(image,dtype=np.float32)/255)[...,0].astype(np.float32) |
| h,w=L.shape;scale=min(size/max(h,w),1) |
| target=(max(8,round(w*scale)),max(8,round(h*scale))) |
| small=np.asarray(Image.fromarray(L).resize(target,Image.Resampling.BILINEAR),dtype=np.float32) |
| x=(small[None,None]/50-1).copy() |
| ab=session.run(['chroma'],{'luminance':x})[0][0] |
| full=np.stack([np.asarray(Image.fromarray(c).resize((w,h),Image.Resampling.BILINEAR)) for c in ab],axis=-1) |
| rgb=np.clip(lab2rgb(np.concatenate([L[...,None],full],axis=-1)),0,1) |
| return Image.fromarray(np.rint(rgb*255).astype(np.uint8)) |
|
|
| if __name__=='__main__': |
| p=argparse.ArgumentParser(__doc__);p.add_argument('images',nargs='+');p.add_argument('--model',required=True) |
| p.add_argument('--output-dir',default='colorized');p.add_argument('--size',type=int,default=256);p.add_argument('--threads',type=int,default=2) |
| a=p.parse_args();session=load_session(a.model,a.threads);out=Path(a.output_dir);out.mkdir(parents=True,exist_ok=True) |
| for source in a.images: |
| with Image.open(source) as im:result=colorize(session,im,a.size) |
| dest=out/(Path(source).stem+'_colorized.png');result.save(dest);print(dest) |
|
|
