cupido/scripts/track_videos.py

"""Headless offline tracker.

Reads target JSONs produced by `pick_targets.py`, builds the 6 ROIs of the
HD mating arena from the L-shape reference points, runs ethoscope's
`MultiFlyTracker` against the merged.mp4 file via `MovieVirtualCamera`, and
writes a SQLite DB to `data/tracked/<video_basename>_tracking.db`.

Idempotent: skips videos whose tracking DB already exists (unless --redo).

Usage:
    python track_videos.py                # process all videos with target JSON
    python track_videos.py --redo         # re-track even if DB exists
    python track_videos.py --jobs 4       # run up to 4 videos in parallel
    python track_videos.py --max-duration 1800  # cap each video at 30 min (sec)
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import sys
import traceback
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

import numpy as np

# Import ethoscope from the local source tree (no pip install).
ETHOSCOPE_SRC = Path("/home/gg/Code/ethoscope_project/ethoscope/src/ethoscope")
sys.path.insert(0, str(ETHOSCOPE_SRC))

from config import TARGETS_DIR, TRACKING_OUTPUT_DIR  # noqa: E402
from tracking_geometry import HD_FG_DATA, compute_roi_polygons  # noqa: E402


def build_rois_from_targets(reference_points):
    """Wrap the shared geometry into ethoscope `ROI` objects."""
    from ethoscope.core.roi import ROI

    polys = compute_roi_polygons(reference_points)
    return [ROI(poly.reshape((1, 4, 2)), idx=i + 1) for i, poly in enumerate(polys)]


def track_one(json_path: Path, output_dir: Path, max_duration: float | None,
              redo: bool) -> tuple[str, str]:
    """Track a single video. Returns (status, message). Run in subprocess.

    Statuses: "ok", "skip", "error".
    """
    # Re-import inside subprocess so each worker has its own ethoscope state.
    import sys as _sys
    _sys.path.insert(0, str(ETHOSCOPE_SRC))
    import cv2
    from ethoscope.core.monitor import Monitor
    from ethoscope.hardware.input.cameras import MovieVirtualCamera
    from ethoscope.io.sqlite import SQLiteResultWriter
    from ethoscope.trackers.multi_fly_tracker import MultiFlyTracker

    class BGRMovieCamera(MovieVirtualCamera):
        """MovieVirtualCamera variant that keeps BGR frames.

        MultiFlyTracker calls cv2.cvtColor(img, COLOR_BGR2GRAY) without checking
        whether img is already grayscale, so we must feed it 3-channel input.
        """
        def _next_image(self):
            ret, frame = self.capture.read()
            if not ret or frame is None:
                return None
            return frame  # BGR, untouched

    payload = json.loads(json_path.read_text())
    if payload.get("unusable"):
        reason = payload.get("reason") or "no reason given"
        return "skip", f"marked unusable: {reason}"
    video_path = Path(payload["video_path"])
    if not video_path.exists():
        return "error", f"video missing: {video_path}"

    out_db = output_dir / f"{video_path.stem}_tracking.db"
    if out_db.exists() and not redo:
        return "skip", f"DB exists: {out_db.name}"
    if out_db.exists():
        out_db.unlink()

    rois = build_rois_from_targets(payload["reference_points"])

    cam_kwargs = {"use_wall_clock": False}
    if max_duration is not None:
        cam_kwargs["max_duration"] = max_duration
    cam = BGRMovieCamera(str(video_path), **cam_kwargs)

    metadata = {
        "machine_id": payload.get("machine_uuid", "unknown"),
        "machine_name": payload.get("machine_name", "unknown"),
        "date_time": int(payload.get("session_epoch", 0)),
        "frame_width": cam.width,
        "frame_height": cam.height,
        "version": "offline-tracker-1",
        "experimental_info": "{}",
        "selected_options": json.dumps({
            "tracker": "MultiFlyTracker",
            "template": "HD_Mating_Arena_6_ROIS",
            "fg_data": HD_FG_DATA,
            "maxN": 2,
        }),
        "hardware_info": "{}",
        "reference_points": str([list(map(int, p)) for p in payload["reference_points"]]),
        "backup_filename": out_db.name,
        "result_writer_type": "SQLite3",
        "sqlite_source_path": str(out_db),
    }

    tracker_data = {
        "maxN": 2,
        "visualise": False,
        "fg_data": HD_FG_DATA,
        "adaptive_threshold": True,
        "min_fg_threshold": 10,
        "max_fg_threshold": 50,
    }

    db_credentials = {"name": str(out_db)}
    rw = SQLiteResultWriter(
        db_credentials, rois, metadata=metadata,
        make_dam_like_table=False, take_frame_shots=False, erase_old_db=True,
    )

    monit = Monitor(
        cam, MultiFlyTracker, rois,
        reference_points=payload["reference_points"],
        data=tracker_data,
    )

    try:
        with rw as result_writer:
            monit.run(result_writer=result_writer, drawer=None, verbose=False)
    except Exception:
        return "error", traceback.format_exc(limit=5)
    finally:
        try:
            cam._close()
        except Exception:
            pass

    if not out_db.exists():
        return "error", "tracking finished but DB was not created"
    return "ok", str(out_db)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--redo", action="store_true", help="re-track even if DB exists")
    parser.add_argument("--jobs", type=int, default=1, help="parallel workers")
    parser.add_argument(
        "--max-duration", type=float, default=None,
        help="cap each video at this many seconds (default: full video)",
    )
    parser.add_argument("--limit", type=int, default=None, help="process only first N")
    parser.add_argument("--video", type=str, default=None,
                        help="track a single video (mp4 path); requires its target JSON")
    args = parser.parse_args()

    TRACKING_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    if args.video:
        stem = Path(args.video).stem
        json_path = TARGETS_DIR / f"{stem}.json"
        if not json_path.exists():
            sys.exit(f"No target JSON for {args.video}: expected {json_path}")
        jsons = [json_path]
    else:
        jsons = sorted(TARGETS_DIR.glob("*.json"))

    if args.limit:
        jsons = jsons[: args.limit]

    if not jsons:
        print("No target JSONs found. Run pick_targets.py first.")
        return

    print(f"Tracking {len(jsons)} videos (jobs={args.jobs}, redo={args.redo}).")
    n_ok = n_skip = n_err = 0

    if args.jobs <= 1:
        for jp in jsons:
            print(f"  → {jp.name}", flush=True)
            status, msg = track_one(jp, TRACKING_OUTPUT_DIR, args.max_duration, args.redo)
            print(f"    {status}: {msg.splitlines()[-1] if msg else ''}", flush=True)
            n_ok += status == "ok"
            n_skip += status == "skip"
            n_err += status == "error"
    else:
        with ProcessPoolExecutor(max_workers=args.jobs) as ex:
            futs = {
                ex.submit(track_one, jp, TRACKING_OUTPUT_DIR, args.max_duration, args.redo): jp
                for jp in jsons
            }
            for fut in as_completed(futs):
                jp = futs[fut]
                try:
                    status, msg = fut.result()
                except Exception as e:
                    status, msg = "error", f"future raised: {e}"
                print(f"  {jp.name}: {status} — {msg.splitlines()[-1] if msg else ''}",
                      flush=True)
                n_ok += status == "ok"
                n_skip += status == "skip"
                n_err += status == "error"

    print(f"\nDone. ok={n_ok}  skipped={n_skip}  errors={n_err}")
    sys.exit(0 if n_err == 0 else 1)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
    main()