#!/usr/bin/env python3
"""Cut the 3D kid from the flat grey studio background."""
import numpy as np, onnxruntime as ort
from PIL import Image, ImageFilter
from scipy import ndimage

SRC = "/root/notchlight/assets3d/kid_flashlight.png"
OUT = "/root/notchlight/assets3d/kid_cut.png"

img = Image.open(SRC).convert("RGB")
sess = ort.InferenceSession("/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], np.float32)) / np.array([0.229, 0.224, 0.225], np.float32)
x = np.transpose(x, (2, 0, 1))[None].astype(np.float32)
pred = sess.run(None, {name: x})[0][0, 0]
pred = (pred - pred.min()) / (pred.max() - pred.min())
mask = Image.fromarray((pred * 255).astype(np.uint8)).resize(img.size, Image.LANCZOS)

a = np.asarray(img).astype(np.int16)
mx = a.max(axis=2); mn = a.min(axis=2)
sat = mx - mn
# the studio grey is desaturated; the kid's warm palette is not.
greyish = (sat < 14) & (mx > 70)
# only let greyish pixels outside the confident u2net core become transparent
core = np.asarray(mask) > 210
alpha = np.asarray(mask).copy()
alpha[greyish & ~ndimage.binary_erosion(core, iterations=3)] = 0

# clean up speckle and soften the edge
alpha = ndimage.binary_opening(alpha > 100, iterations=1) * 255
alpha = np.asarray(Image.fromarray(alpha.astype(np.uint8)).filter(ImageFilter.GaussianBlur(0.7)))

rgba = np.dstack([a.astype(np.uint8), alpha.astype(np.uint8)])
kid = Image.fromarray(rgba, "RGBA")
ys, xs = np.where(alpha > 8)
kid = kid.crop((xs.min(), ys.min(), xs.max() + 1, ys.max() + 1))
kid.save(OUT)
print("kid cut", kid.size)
bg = Image.new("RGB", kid.size, (10, 14, 28)); bg.paste(kid, (0, 0), kid); bg.save("/tmp/kid3d_dark.png")