# Clips freistellen: python3 werkzeuge/clip_freistellen.py <arbeitsordner> <zielordner> clip1.mp4 ...
# Clip freistellen: Frames -> Maske -> gemeinsamer Zuschnitt -> WebM (VP9 mit Alpha)
import sys, os, subprocess, json, glob
import numpy as np
from PIL import Image, ImageFilter
from scipy import ndimage
from rembg import remove, new_session
s = new_session("isnet-general-use")
work, ziel = sys.argv[1], sys.argv[2]
os.makedirs(ziel, exist_ok=True)
for src in sys.argv[3:]:
    name = os.path.basename(src)[:-4]
    if os.path.exists(f"{ziel}/{name}.webm"): continue
    d = f"{work}/{name}"; os.makedirs(d + "/roh", exist_ok=True); os.makedirs(d + "/frei", exist_ok=True)
    subprocess.run(["ffmpeg", "-v", "error", "-y", "-i", src, f"{d}/roh/%04d.png"], check=True)
    roh = sorted(glob.glob(f"{d}/roh/*.png")); boxen = []
    for f in roh:
        dst = f"{d}/frei/" + os.path.basename(f)
        im = Image.open(f).convert("RGB")
        if not os.path.exists(dst):
            m = np.asarray(remove(im, session=s, only_mask=True)).astype(float)
            hard = m > 110
            lab, n = ndimage.label(hard)
            if n:
                sizes = ndimage.sum(hard, lab, range(1, n + 1)); keep = np.zeros_like(hard)
                for i, sz in enumerate(sizes):
                    if sz > 0.02 * sizes.max(): keep |= lab == i + 1
                hard = keep
            holes = ndimage.binary_fill_holes(hard) & ~hard
            hl, hn = ndimage.label(holes)
            for i in range(1, hn + 1):
                if (hl == i).sum() < 400: hard |= hl == i
            hard = ndimage.binary_erosion(hard, iterations=2)
            a = Image.fromarray((hard * 255).astype("uint8")).filter(ImageFilter.GaussianBlur(0.8))
            o = im.convert("RGBA"); o.putalpha(a); o.save(dst)
        boxen.append(Image.open(dst).getchannel("A").getbbox())
    bx = [b for b in boxen if b]
    box = [min(b[0] for b in bx) - 4, min(b[1] for b in bx) - 4, max(b[2] for b in bx) + 4, max(b[3] for b in bx) + 4]
    box = [max(0, box[0]), max(0, box[1]), min(1284, box[2]), min(716, box[3])]
    if (box[2] - box[0]) % 2: box[2] -= 1
    if (box[3] - box[1]) % 2: box[3] -= 1
    w, h = box[2] - box[0], box[3] - box[1]
    subprocess.run(["ffmpeg", "-v", "error", "-y", "-framerate", "24", "-i", f"{d}/frei/%04d.png",
        "-vf", f"crop={w}:{h}:{box[0]}:{box[1]}", "-c:v", "libvpx-vp9", "-pix_fmt", "yuva420p",
        "-b:v", "0", "-crf", "28", "-auto-alt-ref", "0", f"{ziel}/{name}.webm"], check=True)
    json.dump({"box": box, "boxen": boxen, "frames": len(roh)}, open(f"{ziel}/{name}.json", "w"))
    print(name, box, len(roh), flush=True)