"""Stage 1 of the semantic pipeline: for each archive clip, pull frames, estimate depth (MiDaS v2.1 small),
find people (U2-Net human segmentation) and salient objects (U2-Net-P), and label every pixel with one of five classes:
SOURCE (things that emit light), PERSON, OBJECT, SKY, WORLD. Writes one npz per clip: rgb, inverse depth, labels, per frame.
python3 prep.py key [key ...]"""
import sys, os, json, subprocess, numpy as np, cv2, onnxruntime as ort
A = '/home/claude/afterlight/archive/'; M = '/home/claude/afterlight/semantic/models/'; OUT = '/home/claude/afterlight/semantic/cache/'
os.makedirs(OUT, exist_ok=True)
SI = json.load(open(A + 'sheet_index.json'))
W, H = 480, 270
opts = ort.SessionOptions(); opts.intra_op_num_threads = 2
DEP = ort.InferenceSession(M + 'midas_small.onnx', opts); SAL = ort.InferenceSession(M + 'u2netp.onnx', opts); HUM = ort.InferenceSession(M + 'u2net_human_seg.onnx', opts)
MEAN, STD = np.array([0.485, 0.456, 0.406], np.float32), np.array([0.229, 0.224, 0.225], np.float32)
def prep(img, size):
    x = cv2.resize(img, (size, size), interpolation=cv2.INTER_AREA).astype(np.float32) / 255
    return ((x - MEAN) / STD).transpose(2, 0, 1)[None].astype(np.float32)
def depth_raw(img):
    d = DEP.run(None, {DEP.get_inputs()[0].name: prep(img, 256)})[0][0]
    return cv2.resize(d, (W, H), interpolation=cv2.INTER_CUBIC)
def depth(img):
    d = depth_raw(img); d = (d - d.min()) / (d.max() - d.min() + 1e-6)
    return d                                   # inverse depth: 1 near, 0 far
def mask(sess, img):
    m = sess.run(None, {sess.get_inputs()[0].name: prep(img, 320)})[0][0, 0]
    m = (m - m.min()) / (m.max() - m.min() + 1e-6); return cv2.resize(m, (W, H))
CLASSES = ['SOURCE', 'PERSON', 'OBJECT', 'SKY', 'WORLD']
def labels(img, inv, sal, hum):
    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV).astype(np.float32); v = hsv[..., 2] / 255; s_ = hsv[..., 1] / 255
    lum = img.astype(np.float32).mean(2) / 255
    grad = cv2.GaussianBlur(np.abs(cv2.Laplacian(cv2.GaussianBlur(lum, (0, 0), 1), cv2.CV_32F)), (0, 0), 6)
    yy = np.linspace(0, 1, H)[:, None] * np.ones((1, W))
    L = np.full((H, W), 4, np.uint8)
    sky = (yy < 0.6) & (inv < 0.25) & (grad < 0.02) & (lum > 0.3)
    L[sky] = 3
    L[(sal > 0.5)] = 2
    L[(hum > 0.6) & (lum < 0.7)] = 1
    thr = max(0.7, float(np.percentile(lum, 94)))
    src = (lum >= thr) & (v > 0.72)
    L[cv2.dilate(src.astype(np.uint8), np.ones((3, 3), np.uint8)) > 0] = 0
    return L
def frames(key, t0, t1, fps):
    path = A + 'clips/' + SI[key] + '.mp4'
    raw = subprocess.run(['ffmpeg', '-v', 'error', '-ss', str(t0), '-t', str(t1 - t0), '-i', path, '-vf', f'fps={fps},scale={W}:{H}:force_original_aspect_ratio=increase,crop={W}:{H}',
                          '-f', 'rawvideo', '-pix_fmt', 'rgb24', '-'], capture_output=True).stdout
    return np.frombuffer(raw, np.uint8).reshape(-1, H, W, 3)
def run(key, t0=0.3, t1=None, fps=24, secs=6.0):
    """Full frame rate. Depth on every frame, normalized with a slowly moving range and smoothed in time so the cloud does not flicker;
    the two segmentation nets (slow) run every sixth frame and are held between."""
    path = A + 'clips/' + SI[key] + '.mp4'
    dur = float(subprocess.run(['ffprobe', '-v', 'error', '-show_entries', 'format=duration', '-of', 'csv=p=0', path], capture_output=True, text=True).stdout)
    t1 = min(dur - 0.1, t1 or t0 + secs)
    F = frames(key, t0, t1, fps); rgb, inv, lab = [], [], []
    lo = hi = None; prev = None
    for k, img in enumerate(F):
        r = depth_raw(img); a, b = np.percentile(r, 1), np.percentile(r, 99)
        lo = a if lo is None else 0.85 * lo + 0.15 * a; hi = b if hi is None else 0.85 * hi + 0.15 * b
        d = np.clip((r - lo) / (hi - lo + 1e-6), 0, 1)
        d = d if prev is None else 0.55 * d + 0.45 * prev; prev = d
        if k % 6 == 0: s = mask(SAL, img); h = mask(HUM, img)
        rgb.append(img); inv.append(d.astype(np.float16)); lab.append(labels(img, d, s, h))
    np.savez_compressed(OUT + key + '.npz', rgb=np.array(rgb), inv=np.array(inv), lab=np.array(lab), fps=fps)
    counts = np.bincount(np.array(lab).ravel(), minlength=5) / np.array(lab).size
    print(key, len(F), 'frames', {c: round(float(p_), 3) for c, p_ in zip(CLASSES, counts)}, flush=True)
if __name__ == '__main__':
    for k in sys.argv[1:]: run(k)