video_processor.py

"""
Video processor — sampled keyframe inference with multi-object tracking.
Eliminates double-counting by tracking unique persons/cars across frames.
"""
import time
import base64
import cv2
import numpy as np

from inference import CrowdCounter
from tracker import CentroidTracker


class VideoProcessor:
    def __init__(self, counter: CrowdCounter, sample_interval: int = 30,
                 track_distance: float = 120):
        self.counter = counter
        self.sample_interval = sample_interval
        self.tracker = CentroidTracker(
            max_disappeared=3,
            min_appeared=2,
            distance_threshold=track_distance,
        )

    def process(self, video_path: str) -> dict:
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise ValueError(f"Cannot open video: {video_path}")

        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_counts = []
        frame_idx = 0
        t0 = time.time()

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            if frame_idx % self.sample_interval == 0:
                result = self.counter.count(frame)
                detections = result.get("detections", [])

                # Update tracker with this frame's detections
                track_result = self.tracker.update(detections, frame_idx)

                heatmap_b64 = base64.b64encode(result["heatmap_bytes"]).decode("utf-8")
                vis_b64 = base64.b64encode(result["vis_bytes"]).decode("utf-8")

                frame_counts.append({
                    "frame": frame_idx,
                    "count": result["count_rounded"],
                    "count_raw": round(result["count"], 2),
                    "person_count": result["person_count"],
                    "car_count": result["car_count"],
                    "heatmap_base64": heatmap_b64,
                    "vis_base64": vis_b64,
                    "active_tracks": track_result["active_tracks"],
                    "unique_counts": track_result["unique_counts"],
                    "processing_time_ms": result["processing_time_ms"],
                })

            frame_idx += 1

        cap.release()
        elapsed = (time.time() - t0) * 1000

        # Final tracking summary
        tracking_summary = self.tracker.final_summary()

        counts = [fc["count"] for fc in frame_counts]
        return {
            "frame_counts": frame_counts,
            "average_count": round(float(np.mean(counts)), 1) if counts else 0,
            "max_count": int(np.max(counts)) if counts else 0,
            "min_count": int(np.min(counts)) if counts else 0,
            "total_frames": total_frames,
            "fps": round(fps, 1),
            "sampled_frames": len(frame_counts),
            "processing_time_ms": round(elapsed, 1),
            "tracking_summary": tracking_summary,
            # Unique counts from tracker (not per-frame sums)
            "unique_person_count": tracking_summary.get("person", {}).get("count", 0),
            "unique_car_count": tracking_summary.get("car", {}).get("count", 0),
        }