#!/usr/bin/env python3
"""Cutout v5: rembg 2x + white-bg decontamination + heavy smooth alpha. No pixel fringes."""
import numpy as np
from PIL import Image, ImageFilter
from rembg import remove, new_session
import gc

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 decontaminate_white(rgb, alpha):
    """Remove white fringe: c = (rgb - white*(1-a))/a, clamped."""
    a = np.clip(alpha, 0.0, 1.0)[..., None]
    out = (rgb.astype(np.float32) - 255.0 * (1.0 - a)) / np.maximum(a, 1e-6)
    return np.clip(out, 0, 255).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)
        # supersample: upscale the segmented result back to 3x for smooth edges
        rgba = rgba.resize((int(cell.width*3), int(cell.height*3)), Image.LANCZOS)
        gc.collect()
        a0 = np.array(rgba)[:, :, 3].astype(np.float32) / 255.0
        rgb0 = np.array(rgba)[:, :, :3]

        # decontaminate from white background
        rgb_dc = decontaminate_white(rgb0, a0)

        if key in light:
            # dilate slightly + heavy blur -> soft feathered edge, no staircase
            a_im = Image.fromarray((a0*255).astype(np.uint8), 'L').filter(ImageFilter.MaxFilter(5))
            a_im = a_im.filter(ImageFilter.GaussianBlur(2.8))
            a2 = np.array(a_im).astype(np.float32) / 255.0
        else:
            a_im = Image.fromarray((a0*255).astype(np.uint8), 'L').filter(ImageFilter.MinFilter(7))
            a_im = a_im.filter(ImageFilter.GaussianBlur(2.8))
            a2 = np.array(a_im).astype(np.float32) / 255.0

        # final rgb: decontaminated where opaque-ish, original where transparent
        mix = np.clip(a2 / np.maximum(a0, 1e-6), 0, 1)[..., None]
        rgb_final = (rgb_dc * mix + rgb0 * (1 - mix)).astype(np.uint8)

        out = Image.fromarray(np.dstack([rgb_final, (a2*255).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))
        del rgba, a0, rgb0, rgb_dc, out
        gc.collect()
print('done')