| |
| import importlib.util,json |
| from pathlib import Path |
| import os |
| import cv2,numpy as np,torch |
| ROOT=Path(os.environ.get('BRAINMU_WORKDIR',Path(__file__).resolve().parents[1])).resolve() |
| spec=importlib.util.spec_from_file_location('tr',Path(__file__).with_name('train_recon.py'));tr=importlib.util.module_from_spec(spec);spec.loader.exec_module(tr) |
| model=tr.Net().cuda().eval();metrics={};ck=torch.load(ROOT/'runs/recon_base_v1/best.pt',map_location='cpu',weights_only=False);model.load_state_dict(ck['model']) |
| with torch.no_grad(): |
| for split in ['train','val','test']: |
| recs=json.loads((ROOT/'artifacts'/f'{split}_pairs.json').read_text());out=ROOT/'data/basenet'/split |
| for d in ['input_basenet','target_clean']: (out/d).mkdir(parents=True,exist_ok=True) |
| manifest=[] |
| for i,r in enumerate(recs): |
| x=torch.from_numpy(tr.load_dat(r['spike']))[None].cuda();p=model(x).clamp(0,1).float().cpu().numpy()[0,0];img=np.round(p*255).astype(np.uint8);name=r['id']+'.png';cv2.imwrite(str(out/'input_basenet'/name),img);rgb=cv2.imread(r['gt_rgb']);cv2.imwrite(str(out/'target_clean'/name),rgb) |
| manifest.append({'id':r['id'],'scene':r['scene'],'split':split,'input':f'input_basenet/{name}','target':f'target_clean/{name}','prompt':'Restore the clean image.','source_dat':r['spike'],'source_gt':r['gt_rgb']}) |
| if (i+1)%100==0:print(split,i+1,flush=True) |
| (out/'manifest.jsonl').write_text(''.join(json.dumps(x)+'\n' for x in manifest));metrics[split]=tr.evaluate(model,split);print('EXPORT_COMPLETE',split,len(manifest),metrics[split]['psnr_mean_db'],metrics[split]['ssim_mean'],flush=True) |
| (ROOT/'artifacts/basenet_metrics.json').write_text(json.dumps(metrics,indent=2)+'\n') |
|
|
