#!/usr/bin/env python3
"""Final cutout pipeline: rembg alpha + per-item erode/dilate + interior color fill (kills halos)."""
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)]

# light items (milky/white tees & hoodie) need dilation; black items need erosion
light = {'r1c1', 'r1c3', 'r3c3'}

def color_fill_boundary(rgb, alpha, fill_iters=12):
    """Replace boundary pixel colors with nearest interior color (kills white halos)."""
    h, w = alpha.shape
    # interior = eroded alpha
    a_im = Image.fromarray(alpha, 'L').filter(ImageFilter.MinFilter(9))
    interior = np.array(a_im) > 200
    out = rgb.copy()
    # mask of pixels needing fill: opaque but not interior
    need = (alpha > 40) & (~interior)
    done = interior.copy()
    # iterative propagation of interior colors into boundary
    for _ in range(fill_iters):
        if not need.any():
            break
        # for each unfilled pixel adjacent to a done pixel, take avg color of done neighbors
        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] += rgb[sy, sx] * d[..., None]
        fill_here = need & (nbr > 0)
        if not fill_here.any():
            break
        out[fill_here] = (colors[fill_here] / nbr[fill_here, None]).astype(np.uint8)
        done |= fill_here
        need &= ~fill_here
    return out

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)
        a = np.array(rgba)[:, :, 3]
        rgb = np.array(rgba)[:, :, :3].astype(np.float32)

        if key in light:
            # dilate to restore chewed edges, mild
            a_im = Image.fromarray(a, 'L').filter(ImageFilter.MaxFilter(5)).filter(ImageFilter.GaussianBlur(0.9))
            a2 = np.array(a_im).astype(np.float32)
        else:
            # erode to kill halo
            a_im = Image.fromarray(a, 'L').filter(ImageFilter.MinFilter(7)).filter(ImageFilter.GaussianBlur(0.9))
            a2 = np.array(a_im).astype(np.float32)

        rgb_filled = color_fill_boundary(rgb.astype(np.uint8), np.array(a).astype(np.uint8))
        # soft edge: blend raw rgb into filled rgb by alpha coverage (keeps texture, drops halo)
        a_norm = (a2[..., None] / 255.0)
        rgb_final = (rgb_filled * a_norm + rgb * (1 - a_norm)).astype(np.uint8)
        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')