#!/usr/bin/env python3
"""Cut the kid out with U2Net and prepare the 1206x2622 background plate."""
import numpy as np, onnxruntime as ort
from PIL import Image, ImageFilter

ROOT = "/root/notchlight"

# ---------- background: crop the cream print border, scale to exact screen ----------
bg = Image.open(f"{ROOT}/assets/bg_v1.png").convert("RGB")
# inner artwork box measured from the cream margin, plus a small inset past the dark frame line
box = (18 + 9, 64 + 9, 1138 - 9, 2540 - 9)
art = bg.crop(box)
target = (1206, 2622)
scale = max(target[0] / art.width, target[1] / art.height)
art = art.resize((round(art.width * scale), round(art.height * scale)), Image.LANCZOS)
left = (art.width - target[0]) // 2
top = 0  # keep the floor, crop empty sky from the top
art = art.crop((left, top, left + target[0], top + target[1]))
art.save(f"{ROOT}/assets/bg_plate.png")
print("bg plate", art.size)

# ---------- kid: u2net alpha ----------
img = Image.open(f"{ROOT}/assets/kid_v2.png").convert("RGB")
sess = ort.InferenceSession(f"{ROOT}/../.u2net/u2net.onnx", providers=["CPUExecutionProvider"])
name = sess.get_inputs()[0].name

inp = img.resize((320, 320), Image.LANCZOS)
x = np.asarray(inp, dtype=np.float32) / 255.0
x = (x - np.array([0.485, 0.456, 0.406], dtype=np.float32)) / np.array([0.229, 0.224, 0.225], dtype=np.float32)
x = np.transpose(x, (2, 0, 1))[None].astype(np.float32)
out = sess.run(None, {name: x})[0][0, 0]
out = (out - out.min()) / (out.max() - out.min())
mask = Image.fromarray((out * 255).astype(np.uint8)).resize(img.size, Image.LANCZOS)

# tighten the edge very slightly, then feather
mask = mask.filter(ImageFilter.MinFilter(3)).filter(ImageFilter.GaussianBlur(0.8))
arr = np.asarray(img).astype(np.float32)
alpha = np.asarray(mask).astype(np.float32) / 255.0
# decontaminate the cream fringe on semi-transparent edge pixels
bgc = np.array([244, 236, 218], dtype=np.float32)
a3 = alpha[..., None]
rgb = np.where(a3 > 0.02, (arr - (1 - a3) * bgc) / np.maximum(a3, 1e-3), arr)
rgb = np.clip(rgb, 0, 255)

rgba = np.dstack([rgb, alpha * 255]).astype(np.uint8)
kid = Image.fromarray(rgba, "RGBA")
kid = kid.crop(kid.getbbox())
kid.save(f"{ROOT}/assets/kid_cut.png")
print("kid cut", kid.size)