#!/usr/bin/env python3
"""Final v2: black items = erode + full interior color fill (no raw re-blend). Light items = dilate only."""
import numpy as np
from PIL import Image, ImageFilter
from rembg import remove, new_session

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

im = Image.open(SRC).convert('RGB')
rows = [(265, 611), (731, 1121), (1246, 1608)]
cols = [(27, 471), (524, 977), (1028, 1470)]
light = {'r1c1', 'r1c3', 'r3c3'}

def fill_interior(rgb, alpha, iters=60):
    """Propagate nearest interior color to ALL opaque pixels (halo-free)."""
    h, w = alpha.shape
    a_im = Image.fromarray(alpha, 'L').filter(ImageFilter.MinFilter(13))
    interior = np.array(a_im) > 200
    out = rgb.astype(np.float32).copy()
    done = interior.copy()
    need = ~interior
    for _ in range(iters):
        if not need.any():
            break
        nbr = np.zeros((h, w)); colors = np.zeros((h, w, 3))
        for dy, dx in ((1,0),(-1,0),(0,1),(0,-1),(1,1),(1,-1),(-1,1),(-1,-1)):
            sy = slice(max(0,dy), h+min(0,dy)); sx = slice(max(0,dx), w+min(0,dx))
            ty = slice(max(0,-dy), h+min(0,-dy)); tx = slice(max(0,-dx), w+min(0,-dx))
            d = done[sy, sx]
            nbr[ty, tx] += d
            colors[ty, tx] += out[sy, sx] * d[..., None]
        fh = need & (nbr > 0)
        if not fh.any():
            break
        out[fh] = colors[fh] / nbr[fh, None]
        done |= fh
        need &= ~fh
    return out.astype(np.uint8)

for ri, (y0, y1) in enumerate(rows):
    for ci, (x0, x1) in enumerate(cols):
        key = f'r{ri+1}c{ci+1}'
        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))
        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)
        a0 = np.array(rgba)[:, :, 3]
        rgb0 = np.array(rgba)[:, :, :3]

        if key in light:
            a_im = Image.fromarray(a0, 'L').filter(ImageFilter.MaxFilter(5)).filter(ImageFilter.GaussianBlur(3.0))
            a2 = np.array(a_im).astype(np.float32)
            rgb_final = rgb0
        else:
            a_im = Image.fromarray(a0, 'L').filter(ImageFilter.MinFilter(9)).filter(ImageFilter.GaussianBlur(3.0))
            a2 = np.array(a_im).astype(np.float32)
            rgb_final = fill_interior(rgb0, a0)

        out = Image.fromarray(np.dstack([rgb_final, a2.astype(np.uint8)]), 'RGBA')
        bbox = out.getbbox()
        if bbox:
            xa, ya, xb, yb = bbox
            m = 3
            out = out.crop((max(0,xa-m), max(0,ya-m), min(out.width,xb+m), min(out.height,yb+m)))
        out.save(OUT + f'item_{key}.png')
        op = np.array(out)[:, :, 3]
        print(key, out.size, 'opq%%=%.1f' % ((op>200).mean()*100))
print('done')