"""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)