predict_video.py

"""
Real-time webcam / video file FER inference.

Usage:
    python predict_video.py --source 0               # webcam
    python predict_video.py --source video.mp4       # video file
    python predict_video.py --source 0 --save-output out.mp4
"""

import argparse
import sys
import time
from pathlib import Path

import cv2
import numpy as np

sys.path.insert(0, str(Path(__file__).parent))

from inference import FERPredictor
from utils import EMOTION_BGR


DETECT_EVERY_N = 3   # run face detection every N frames; reuse bbox in between
FACE_INPUT_SCALE = 0.5  # downsample for face detection speed, upsample bbox back


def _scale_bboxes(bboxes, scale):
    """Scale (x,y,w,h) bboxes back from a downsampled frame."""
    return [(int(x / scale), int(y / scale), int(w / scale), int(h / scale))
            for (x, y, w, h) in bboxes]


def run(source, weights, device, save_path, face_method, no_detect):
    predictor = FERPredictor(weights_path=weights, device=device)

    # Open source
    cap_src = 0 if source == '0' else source
    cap = cv2.VideoCapture(cap_src)
    if not cap.isOpened():
        print(f"[ERROR] Cannot open source: {source}")
        sys.exit(1)

    width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    src_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0

    writer = None
    if save_path:
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        writer = cv2.VideoWriter(save_path, fourcc, src_fps, (width, height))
        print(f"[INFO] Saving output to: {save_path}")

    print("[INFO] Press Q to quit.")

    # State kept between frames
    cached_face_results: list[dict] = []
    frame_idx = 0
    fps_counter = 0
    fps_display = 0.0
    t_last = time.time()

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

        run_detect = (frame_idx % DETECT_EVERY_N == 0)

        if no_detect:
            # Single-crop inference on the whole frame
            if run_detect or not cached_face_results:
                rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                from PIL import Image as PILImage
                pil_frame = PILImage.fromarray(rgb_frame)
                result = predictor.predict_image(pil_frame)
                result['bbox'] = None
                cached_face_results = [result]
        else:
            if run_detect:
                # Downsample before detection
                small = cv2.resize(frame, (0, 0), fx=FACE_INPUT_SCALE, fy=FACE_INPUT_SCALE)
                from detect_face import detect_and_crop_faces
                from PIL import Image as PILImage
                import cv2 as _cv2

                pil_small = PILImage.fromarray(_cv2.cvtColor(small, _cv2.COLOR_BGR2RGB))
                faces = detect_and_crop_faces(
                    pil_small,
                    method=face_method,
                    margin=20,
                    device='cuda' if predictor.device.type == 'cuda' else 'cpu'
                )

                new_results = []
                for idx, (crop, bbox) in enumerate(faces):
                    pred = predictor.predict_image(crop)
                    # Scale bbox back to original resolution
                    if bbox is not None:
                        x, y, w, h = bbox
                        bbox = (
                            int(x / FACE_INPUT_SCALE),
                            int(y / FACE_INPUT_SCALE),
                            int(w / FACE_INPUT_SCALE),
                            int(h / FACE_INPUT_SCALE),
                        )
                    pred['bbox'] = bbox
                    pred['face_index'] = idx
                    new_results.append(pred)
                cached_face_results = new_results if new_results else cached_face_results

        # Draw overlays
        for res in cached_face_results:
            bbox = res.get('bbox')
            emotion = res['emotion']
            conf = res['confidence']
            color = EMOTION_BGR.get(emotion, (200, 200, 200))

            if bbox is not None:
                x, y, w, h = bbox
                cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
                label = f"{emotion} {conf*100:.0f}%"
                font = cv2.FONT_HERSHEY_SIMPLEX
                scale_f, thick = 0.65, 2
                (tw, th), bl = cv2.getTextSize(label, font, scale_f, thick)
                ty = max(y - 6, th + 6)
                cv2.rectangle(frame, (x, ty - th - 4), (x + tw + 4, ty + bl), color, cv2.FILLED)
                cv2.putText(frame, label, (x + 2, ty - 2), font, scale_f,
                            (255, 255, 255), thick, cv2.LINE_AA)
            else:
                # No bbox — overlay emotion in top-center
                label = f"{emotion} {conf*100:.0f}%"
                font = cv2.FONT_HERSHEY_SIMPLEX
                cv2.putText(frame, label, (width // 2 - 80, 40), font, 0.9,
                            EMOTION_BGR.get(emotion, (200, 200, 200)), 2, cv2.LINE_AA)

        # FPS counter
        fps_counter += 1
        if fps_counter >= 10:
            t_now = time.time()
            fps_display = fps_counter / (t_now - t_last)
            t_last = t_now
            fps_counter = 0

        cv2.putText(frame, f"FPS: {fps_display:.1f}", (10, 28),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, cv2.LINE_AA)

        cv2.imshow('FER Inference — Q to quit', frame)

        if writer:
            writer.write(frame)

        if cv2.waitKey(1) & 0xFF in (ord('q'), ord('Q'), 27):
            break

        frame_idx += 1

    cap.release()
    if writer:
        writer.release()
    cv2.destroyAllWindows()
    print("[INFO] Done.")


def main():
    parser = argparse.ArgumentParser(description='Real-time FER on webcam or video.')
    parser.add_argument('--source', default='0',
                        help='Webcam index (0) or path to video file.')
    parser.add_argument('--weights', default='../models/model_weights.pth',
                        help='Path to model_weights.pth.')
    parser.add_argument('--device', default='auto',
                        help='Device: auto | cpu | cuda.')
    parser.add_argument('--save-output', default=None,
                        help='Save annotated video to this path.')
    parser.add_argument('--face-method', default='mtcnn', choices=['mtcnn', 'haar'],
                        help='Face detection backend (default: mtcnn).')
    parser.add_argument('--no-detect', action='store_true',
                        help='Skip face detection — infer on full frame.')

    args = parser.parse_args()
    run(
        source=args.source,
        weights=args.weights,
        device=args.device,
        save_path=args.save_output,
        face_method=args.face_method,
        no_detect=args.no_detect,
    )


if __name__ == '__main__':
    main()