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.automatic_mask_generator import SAM2AutomaticMaskGenerator
ROOT=Path(__file__).resolve().parent;source=Image.open(ROOT/'flatlay.png').convert('RGB')
model=build_sam2('configs/sam2.1/sam2.1_hiera_t.yaml',a.checkpoint,device=a.device)
generator=SAM2AutomaticMaskGenerator(model,points_per_side=16,pred_iou_thresh=0.8,stability_score_thresh=0.9)
with torch.inference_mode():annotations=generator.generate(np.array(source))
out=ROOT/'output';out.mkdir(exist_ok=True);rows=[]
for i,row in enumerate(annotations):
 mask=Image.fromarray(row['segmentation'].astype('uint8')*255);mask.save(out/f'mask-{i}.png')
 x,y,w,h=map(int,row['bbox']);assert w>0 and h>0
 cutout=source.copy();cutout.putalpha(mask);cutout.crop((x,y,min(source.width,x+w),min(source.height,y+h))).save(out/f'crop-{i}.png')
 rows.append({k:row[k]for k in ['area','bbox','predicted_iou','stability_score'] }|{'mask':f'mask-{i}.png','crop':f'crop-{i}.png','offset':[x,y],'human_approval':'pending'})
(out/'annotations.json').write_text(json.dumps({'model_run':True,'input_size':list(source.size),'annotations':rows},indent=2)+'\n')