#!/usr/bin/env python3
"""Cut out 9 items from 3x3 grid using rembg (u2net) + white-fringe cleanup."""
import numpy as np
from PIL import Image, ImageFilter
from rembg import remove, new_session

SRC = '/root/.hermes/image_cache/img_5692f1c30ecb.jpg'
OUT = '/root/.hermes/image_cache/zrd_cutouts/'
session = new_session('isnet-general-use')

im = Image.open(SRC).convert('RGB')
rows = [(181, 416), (499, 765), (850, 1098)]
cols = [(19, 321), (357, 667), (702, 1003)]

for ri, (y0, y1) in enumerate(rows):
    for ci, (x0, x1) in enumerate(cols):
        pad = 10
        bx0, by0 = max(0, x0-pad), max(0, y0-pad)
        bx1, by1 = min(im.width, x1+pad), min(im.height, y1+pad)
        cell = im.crop((bx0, by0, bx1, by1))
        # upscale small cells for better segmentation
        scale = 2.0
        cell_big = cell.resize((int(cell.width*scale), int(cell.height*scale)), Image.LANCZOS)
        rgba = remove(cell_big, session=session, post_process_mask=True)
        rgba = rgba.resize(cell.size, Image.LANCZOS)
        a = np.array(rgba)[:, :, 3].astype(np.float32)
        # remove faint white halo: pixels that are near-white AND low alpha -> drop
        rgb = np.array(cell).astype(np.float32)
        whiteish = (rgb[:,:,0]>200)&(rgb[:,:,1]>200)&(rgb[:,:,2]>200)
        a[whiteish & (a<120)] = 0
        a = Image.fromarray(a.astype(np.uint8), 'L').filter(ImageFilter.GaussianBlur(0.8))
        rgba = cell.convert('RGBA')
        rgba.putalpha(a)
        bbox = rgba.getbbox()
        if bbox:
            xa, ya, xb, yb = bbox
            m = 4
            rgba = rgba.crop((max(0,xa-m), max(0,ya-m), min(rgba.width,xb+m), min(rgba.height,yb+m)))
        name = f'item_r{ri+1}c{ci+1}.png'
        rgba.save(OUT + name)
        op = np.array(rgba)[:,:,3]
        print(name, rgba.size, 'opaque%%=%.1f' % ((op>200).mean()*100))
print('done')