File size: 2,657 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 | """Export the colorizer AND luminance-guided chroma decoder as one ONNX graph.
Input: N,1,H,W normalized Lab luminance (L*/50-1), minimum8px per side.
Output: N,2,H,W Lab a,b values. Resize/preserve original L* in the app.
"""
import argparse,json
from pathlib import Path
import numpy as np
import torch
from model import load_model
from spatial import guided_chroma
class ChromaPipeline(torch.nn.Module):
def __init__(self,model,radius=8,temperature=.38):
super().__init__();self.model=model;self.radius=radius;self.temperature=temperature
def forward(self,L):
return guided_chroma(L,self.model.decode(self.model(L),self.temperature),self.radius)
def export(a):
import onnx,onnxruntime as ort
torch.set_num_threads(2);torch.manual_seed(2026)
model=ChromaPipeline(load_model(a.model),a.radius).eval()
out=Path(a.output);out.parent.mkdir(parents=True,exist_ok=True)
with torch.inference_mode():
torch.onnx.export(model,torch.zeros(1,1,256,256),str(out),opset_version=17,
input_names=['luminance'],output_names=['chroma'],dynamo=False,
dynamic_axes={'luminance':{0:'batch',2:'height',3:'width'},'chroma':{0:'batch',2:'height',3:'width'}})
onnx.checker.check_model(str(out))
opts=ort.SessionOptions();opts.intra_op_num_threads=2;opts.inter_op_num_threads=1
session=ort.InferenceSession(str(out),sess_options=opts,providers=['CPUExecutionProvider'])
checks=[]
for shape in [(1,1,256,256),(1,1,173,241),(2,1,64,80),(1,1,8,9)]:
x=torch.rand(shape)*2-1
with torch.inference_mode():expected=model(x).numpy()
actual=session.run(None,{'luminance':x.numpy()})[0]
error=np.abs(expected-actual)
assert actual.shape==expected.shape and np.isfinite(actual).all()
np.testing.assert_allclose(actual,expected,atol=.01,rtol=.001)
checks.append({'shape':list(shape),'max_ab_difference':float(error.max()),'mean_ab_difference':float(error.mean())})
result={'input':'N,1,H,W Lab L*/50-1; H,W>=8','output':'N,2,H,W Lab chroma a,b',
'guided_radius':a.radius,'guided_epsilon':.001,'temperature':.38,
'opset':17,'checks':checks,'parameters':sum(p.numel() for p in model.parameters()),
'weights_source':a.model,'runtime':ort.__version__,'bytes':out.stat().st_size}
out.with_suffix('.json').write_text(json.dumps(result,indent=2));print(json.dumps(result,indent=2))
if __name__=='__main__':
p=argparse.ArgumentParser(__doc__);p.add_argument('--model',required=True);p.add_argument('--output',required=True);p.add_argument('--radius',type=int,default=8)
export(p.parse_args())
|