from pathlib import Path
import argparse,hashlib,io,json,time,os
from PIL import Image
from rembg import remove,new_session,__version__
ROOT=Path(__file__).resolve().parent;p=argparse.ArgumentParser();p.add_argument('--only');p.add_argument('--output',default='output');a=p.parse_args()
out=ROOT/a.output;out.mkdir(exist_ok=True);session=new_session('u2netp',providers=['CPUExecutionProvider']);records=[]
for source in sorted((ROOT/'inputs').glob('*')):
 if a.only and source.name!=a.only:continue
 row={'source':source.name,'sha256':hashlib.sha256(source.read_bytes()).hexdigest(),'model':'u2netp'}
 try:
  im=Image.open(source);im.load();result=remove(im.convert('RGB'),session=session)
  temp=out/(source.stem+'.tmp.png');result.save(temp);Image.open(temp).load()
  target=out/(source.stem+'.png');os.replace(temp,target)
  row.update(status='complete',output=target.name,size=list(result.size),mode=result.mode)
 except Exception as e:row.update(status='error',error_type=type(e).__name__,error=str(e).replace(str(ROOT),'[lesson]'))
 records.append(row)
(out/'manifest.json').write_text(json.dumps({'model_run':True,'version':__version__,'records':records},indent=2)+'\n')
print(json.dumps(records))