from pathlib import Path
from PIL import Image, ImageDraw
from rembg import remove, new_session
import json, shutil, numpy as np

root=Path('/Users/bot1/Volumes/root_for_ai/AI工作区/良渚_详情页_一壶二杯青瓷功夫茶具礼盒_20260615_1228')
work=root/'work/main_images_gpt_scene_v05'
cutdir=work/'cutouts'
srcdir=work/'sources'
for p in [work,cutdir,srcdir]: p.mkdir(parents=True,exist_ok=True)
base=root/'work/main_images_from_v03_800x800/assets/images'
sources=[
 ('01','一壶二杯','01_白底套装主KV_重新选图.jpg'),
 ('02','温润出汤','02_倒茶_使用质感.jpg'),
 ('03','茶席成礼','03_俯拍_茶席场景.jpg'),
 ('04','开盒见礼','04_礼盒_开盒礼赠.jpg'),
 ('05','琮纹细节','05_细节_壶流琮纹.jpg'),
]
session=new_session('u2net')
results=[]
for idx,word,fn in sources:
    src=base/fn
    local=srcdir/f'{idx}_{fn}'
    shutil.copy2(src,local)
    im=Image.open(local).convert('RGBA').resize((800,800),Image.Resampling.LANCZOS)
    out=remove(im,session=session,alpha_matting=True,alpha_matting_foreground_threshold=240,alpha_matting_background_threshold=10,alpha_matting_erode_size=8)
    alpha=out.getchannel('A')
    arr=np.array(alpha)
    arr[arr<8]=0
    out.putalpha(Image.fromarray(arr.astype('uint8'),'L'))
    cut=cutdir/f'{idx}_{word}_product_cutout.png'
    out.save(cut)
    bbox=out.getchannel('A').getbbox()
    results.append({'idx':idx,'word':word,'source':str(local),'cutout':str(cut),'bbox':bbox,'alpha_extrema':out.getchannel('A').getextrema()})
cols=len(sources); cell=220
sheet=Image.new('RGB',(cols*cell,cell*2),(245,241,232))
for i,(idx,word,fn) in enumerate(sources):
    orig=Image.open(srcdir/f'{idx}_{fn}').convert('RGB').resize((cell,cell),Image.Resampling.LANCZOS)
    cut=Image.open(cutdir/f'{idx}_{word}_product_cutout.png').convert('RGBA')
    chk=Image.new('RGB',(800,800),(255,255,255)); d=ImageDraw.Draw(chk); s=40
    for y in range(0,800,s):
      for x in range(0,800,s):
        if (x//s+y//s)%2: d.rectangle([x,y,x+s-1,y+s-1],fill=(210,210,210))
    chk.paste(cut,(0,0),cut)
    chk=chk.resize((cell,cell),Image.Resampling.LANCZOS)
    sheet.paste(orig,(i*cell,0)); sheet.paste(chk,(i*cell,cell))
qa=work/'cutout_qa_contact.jpg'; sheet.save(qa,quality=92)
(work/'cutout_manifest.json').write_text(json.dumps(results,ensure_ascii=False,indent=2),encoding='utf-8')
print(json.dumps({'work':str(work),'qa':str(qa),'results':results},ensure_ascii=False,indent=2))