import argparse,json
from pathlib import Path
import numpy as np
import torch
from irstd.models import TinyUNet
from irstd.data import read_manifest,read_pair,sha

def main():
 p=argparse.ArgumentParser();p.add_argument('--manifest',required=True);p.add_argument('--checkpoint',required=True)
 p.add_argument('--out',required=True);p.add_argument('--device',default='cuda');a=p.parse_args()
 ck=torch.load(a.checkpoint,map_location='cpu',weights_only=True)
 model=TinyUNet().to(a.device);model.load_state_dict(ck['state_dict']);model.eval()
 out=Path(a.out);out.mkdir(parents=True,exist_ok=False)
 with torch.no_grad():
  for r in read_manifest(a.manifest):
   x,_=read_pair(r)
   prob=model(torch.from_numpy(x)[None,None].to(a.device)).sigmoid()[0,0].cpu().numpy()
   np.save(out/(r['id']+'.npy'),prob)
 (out/'provenance.json').write_text(json.dumps(dict(checkpoint_sha256=sha(a.checkpoint),manifest_sha256=sha(a.manifest)),indent=2))
if __name__=='__main__':main()