#!/usr/bin/env python3
"""Single-item re-cut with padding + strong erosion (memory-safe)."""
import sys, os
import numpy as np
from PIL import Image, ImageFilter
from rembg import remove, new_session
import gc

session = new_session('isnet-general-use')
OUT = '/root/.hermes/image_cache/zrd_1x1/cutouts/'

def cut(src, out, is_model, erode, blur, pad):
    cell = Image.open(src).convert('RGB')
    if pad:
        c = Image.new('RGB', (cell.width+2*pad, cell.height+2*pad), (255,255,255))
        c.paste(cell, (pad, pad))
        cell = c
    if max(cell.size) > 2000:
        r = 2000 / max(cell.size)
        cell = cell.resize((int(cell.width*r), int(cell.height*r)), Image.LANCZOS)
    scale = 1.2
    cell_big = cell.resize((int(cell.width*scale), int(cell.height*scale)), Image.LANCZOS)
    rgba = remove(cell_big, session=session, post_process_mask=True)
    gc.collect()
    a0 = np.array(rgba)[:, :, 3].astype(np.float32) / 255.0
    rgb0 = np.array(rgba)[:, :, :3]
    a = np.clip(a0, 0.0, 1.0)[..., None]
    rgb_dc = (rgb0.astype(np.float32) - 255.0 * (1.0 - a)) / np.maximum(a, 1e-6)
    rgb_dc = np.clip(rgb_dc, 0, 255).astype(np.uint8)
    a_im = Image.fromarray((a0*255).astype(np.uint8), 'L').filter(ImageFilter.MinFilter(erode))
    a_im = a_im.filter(ImageFilter.GaussianBlur(blur))
    a2 = np.array(a_im).astype(np.float32) / 255.0
    mix = np.clip(a2 / np.maximum(a0, 1e-6), 0, 1)[..., None]
    rgb_final = (rgb_dc * mix + rgb0 * (1 - mix)).astype(np.uint8)
    out_im = Image.fromarray(np.dstack([rgb_final, (a2*255).astype(np.uint8)]), 'RGBA')
    bbox = out_im.getbbox()
    if bbox:
        xa, ya, xb, yb = bbox
        m = 3
        out_im = out_im.crop((max(0,xa-m), max(0,ya-m), min(out_im.width,xb+m), min(out_im.height,yb+m)))
    out_im.save(OUT + out)
    print(out, out_im.size, flush=True)

if __name__ == '__main__':
    import os
    os.makedirs(OUT, exist_ok=True)
    # argv: src out is_model erode blur pad
    src, out, is_model = sys.argv[1], sys.argv[2], int(sys.argv[3])
    erode = int(sys.argv[4]) if len(sys.argv) > 4 else 9
    blur = float(sys.argv[5]) if len(sys.argv) > 5 else 3.0
    pad = int(sys.argv[6]) if len(sys.argv) > 6 else 40
    cut(src, out, is_model, erode, blur, pad)