from pathlib import Path
import argparse,json,sys
p=argparse.ArgumentParser();p.add_argument('--repo',required=True);p.add_argument('--checkpoint',required=True);p.add_argument('--device',default='cuda');a=p.parse_args()
sys.path.insert(0,str(Path(a.repo).resolve()))
import numpy as np,torch
from PIL import Image
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor
ROOT=Path(__file__).resolve().parent;source=Image.open(ROOT/'flatlay.png').convert('RGB');points=json.loads((ROOT/'points.json').read_text())
predictor=SAM2ImagePredictor(build_sam2('configs/sam2.1/sam2.1_hiera_t.yaml',a.checkpoint,device=a.device))
with torch.inference_mode():
 predictor.set_image(np.array(source));masks,scores,_=predictor.predict(point_coords=np.array(points['coords']),point_labels=np.array(points['labels']),multimask_output=True)
out=ROOT/'output';out.mkdir(exist_ok=True);records=[]
for i,(mask,score) in enumerate(zip(masks,scores)):
 alpha=Image.fromarray((mask.astype(np.uint8)*255));alpha.save(out/f'candidate-{i}-mask.png');cutout=source.copy();cutout.putalpha(alpha);cutout.save(out/f'candidate-{i}.png');records.append({'candidate':i,'score':float(score),'mask':f'candidate-{i}-mask.png','human_approval':'pending'})
(out/'candidates.json').write_text(json.dumps({'model_run':True,'checkpoint':a.checkpoint,'points':points,'records':records},indent=2)+'\n')