core/tracking.py

"""Object tracking: follow a player (or anything) across a clip with OpenCV CSRT.

The card/overlay lives in the final 9:16 output space, so we first render the
scene's normalized 9:16 video (same crop/reframe the renderer uses) at a small
size, then run the tracker on THAT. The tracked positions therefore come out
directly as output coordinates (0–1) — no crop-math mapping needed — and they
respect any reframe (pan/zoom) on the clip.
"""

from __future__ import annotations

import os
import tempfile

import cv2

from core.stages.render import VIDEO_FPS, _reframe_vf
from core.utils import log, run_ffmpeg


def track_path(clip_path: str, clip_start: float, dur: float, anchors: list,
               box: dict, reframe: list = None, work_w: int = 360, work_h: int = 640,
               keyframe_every: float = 0.15) -> list[dict]:
    """Track between user ``anchors`` with per-segment drift correction.

    ``anchors`` = sorted ``[{"t","x","y"}]`` (scene seconds + 0–1 position) the user
    placed on the subject. Between each pair we run the tracker, then linearly
    rubber-band its path so it lands exactly on the next anchor — bounding drift,
    which is what makes crowded scenes work. Returns ``[{"t","x","y"}]``.
    """
    anchors = sorted((a for a in (anchors or []) if isinstance(a, dict)), key=lambda a: a["t"])
    if len(anchors) < 2:
        return []
    tmp = tempfile.mktemp(suffix=".mp4")
    vf = _reframe_vf(reframe) + f",scale={work_w}:{work_h}"
    seek = ["-ss", f"{clip_start:.3f}"] if clip_start and clip_start > 0 else []
    try:
        run_ffmpeg(["-stream_loop", "-1", *seek, "-i", clip_path, "-vf", vf,
                    "-t", f"{max(0.3, dur):.3f}", "-an", tmp], label="track-normalize")
    except Exception as e:
        log.warning("tracking: normalize render failed: %s", e)
        return []

    cap = None
    try:
        cap = cv2.VideoCapture(tmp)
        fps = cap.get(cv2.CAP_PROP_FPS) or VIDEO_FPS
        frames = []
        while True:
            ok, fr = cap.read()
            if not ok:
                break
            frames.append(fr)
        if not frames:
            return []
        h, w = frames[0].shape[:2]
        bw = max(12, int(box.get("w", 0.10) * w))
        bh = max(12, int(box.get("h", 0.14) * h))

        keys = [{"t": round(anchors[0]["t"], 3), "x": round(anchors[0]["x"], 4), "y": round(anchors[0]["y"], 4)}]
        for a, b in zip(anchors, anchors[1:]):
            fa = max(0, min(int(a["t"] * fps), len(frames) - 1))
            fb = max(fa, min(int(b["t"] * fps), len(frames) - 1))
            bx = int(a["x"] * w - bw / 2)
            by = int(a["y"] * h - bh / 2)
            tracker = cv2.TrackerCSRT_create()
            tracker.init(frames[fa], (max(0, min(bx, w - bw)), max(0, min(by, h - bh)), bw, bh))
            raw = []  # (t, x, y) tracked across the segment
            lx, ly = a["x"], a["y"]
            for fi in range(fa + 1, fb + 1):
                found, bb = tracker.update(frames[fi])
                if found:
                    lx = min(max((bb[0] + bb[2] / 2) / w, 0.0), 1.0)
                    ly = min(max((bb[1] + bb[3] / 2) / h, 0.0), 1.0)
                raw.append((fi / fps, lx, ly))
            # Drift-correct: rubber-band the path so it ends exactly on anchor b.
            span = max(1e-3, b["t"] - a["t"])
            ex, ey = b["x"] - lx, b["y"] - ly
            last_emit = a["t"]
            for (t, x, y) in raw:
                cx = min(max(x + ex * (t - a["t"]) / span, 0.0), 1.0)
                cy = min(max(y + ey * (t - a["t"]) / span, 0.0), 1.0)
                if t - last_emit >= keyframe_every or t >= b["t"] - 1e-3:
                    keys.append({"t": round(t, 3), "x": round(cx, 4), "y": round(cy, 4)})
                    last_emit = t
            keys[-1] = {"t": round(b["t"], 3), "x": round(b["x"], 4), "y": round(b["y"], 4)}  # pin anchor
        return keys
    except Exception as e:
        log.warning("track_path failed: %s", e)
        return []
    finally:
        if cap is not None:
            cap.release()
        try:
            os.remove(tmp)
        except OSError:
            pass


def track_box(clip_path: str, clip_start: float, dur: float, box: dict,
              reframe: list = None, start_offset: float = 0.0,
              work_w: int = 288, work_h: int = 512,
              keyframe_every: float = 0.2) -> list[dict]:
    """Track ``box`` across the scene's visible video; return position keyframes.

    ``box`` = {"x","y","w","h"} as fractions of the 9:16 frame (centre x/y, size
    w/h) at ``start_offset`` seconds into the scene (the frame you placed the card
    on). Returns ``[{"t","x","y"}]`` with t in seconds from ``start_offset`` and
    x/y the tracked centre (0–1). Empty list on failure.
    """
    tmp = tempfile.mktemp(suffix=".mp4")
    vf = _reframe_vf(reframe) + f",scale={work_w}:{work_h}"
    seek_at = max(0.0, clip_start + max(0.0, start_offset))
    track_dur = max(0.3, dur - max(0.0, start_offset))
    seek = ["-ss", f"{seek_at:.3f}"] if seek_at > 0 else []
    try:
        run_ffmpeg(
            ["-stream_loop", "-1", *seek, "-i", clip_path,
             "-vf", vf, "-t", f"{track_dur:.3f}", "-an", tmp],
            label="track-normalize",
        )
    except Exception as e:
        log.warning("tracking: normalize render failed: %s", e)
        return []

    cap = None
    try:
        cap = cv2.VideoCapture(tmp)
        fps = cap.get(cv2.CAP_PROP_FPS) or VIDEO_FPS
        ok, frame = cap.read()
        if not ok:
            return []
        h, w = frame.shape[:2]
        bw = max(12, int(box.get("w", 0.14) * w))
        bh = max(12, int(box.get("h", 0.20) * h))
        bx = int(box["x"] * w - bw / 2)
        by = int(box["y"] * h - bh / 2)
        bx = max(0, min(bx, w - bw))
        by = max(0, min(by, h - bh))
        tracker = cv2.TrackerCSRT_create()
        tracker.init(frame, (bx, by, bw, bh))

        keys = [{"t": 0.0, "x": round(box["x"], 4), "y": round(box["y"], 4)}]
        last_x, last_y = box["x"], box["y"]
        n, last_emit = 0, 0.0
        while True:
            ok, frame = cap.read()
            if not ok:
                break
            n += 1
            t = n / fps
            found, bb = tracker.update(frame)
            if found:
                last_x = min(max((bb[0] + bb[2] / 2) / w, 0.0), 1.0)
                last_y = min(max((bb[1] + bb[3] / 2) / h, 0.0), 1.0)
            if t - last_emit >= keyframe_every:
                keys.append({"t": round(t, 3), "x": round(last_x, 4), "y": round(last_y, 4)})
                last_emit = t
        # always pin a final keyframe at the end so the card holds to the last spot
        if keys[-1]["t"] < track_dur - 0.05:
            keys.append({"t": round(track_dur, 3), "x": round(last_x, 4), "y": round(last_y, 4)})
        return keys
    except Exception as e:
        log.warning("tracking failed: %s", e)
        return []
    finally:
        if cap is not None:
            cap.release()
        try:
            os.remove(tmp)
        except OSError:
            pass