#!/usr/bin/env python3
"""Per-item cutout runner (memory-safe: one item at a time, cap size)."""
import sys, os
sys.path.insert(0, '/root/.hermes/image_cache/zrd_1x1')

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):
    cell = Image.open(src).convert('RGB')
    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.5
    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)
    filt = ImageFilter.MinFilter(5) if is_model else ImageFilter.MinFilter(9)
    blur = 2.2 if is_model else 3.0
    a_im = Image.fromarray((a0*255).astype(np.uint8), 'L').filter(filt)
    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(0/1)
    src, out, is_model = sys.argv[1], sys.argv[2], int(sys.argv[3])
    cut(src, out, is_model)