inference/visual_inference.py

import os
import cv2
import tempfile
import urllib.request
import numpy as np
import torch
import torchvision.transforms as T
from typing import Tuple

# ── ImageNet normalisation (same as test.ipynb) ───────────────────────────────
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]

FRAME_TFMS = T.Compose([
    T.ToPILImage(),
    T.ToTensor(),
    T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])

NUM_FRAMES  = 45
FRAME_SIZE  = 224
FACE_MARGIN = 0.25

# ── Optional MediaPipe face detector ─────────────────────────────────────────
_mp_detector = None

def _load_face_detector(model_cache_dir: str = "/tmp"):
    global _mp_detector
    if _mp_detector is not None:
        return _mp_detector
    try:
        import mediapipe as mp
        from mediapipe.tasks.python import vision as mp_vision
        from mediapipe.tasks.python import BaseOptions

        model_path = os.path.join(model_cache_dir, "blaze_face_full_range.tflite")
        if not os.path.exists(model_path):
            print("Downloading MediaPipe face detector model...")
            urllib.request.urlretrieve(
                "https://storage.googleapis.com/mediapipe-models/face_detector/"
                "blaze_face_full_range/float16/1/blaze_face_full_range.tflite",
                model_path,
            )
            print("  Downloaded ✓")

        options = mp_vision.FaceDetectorOptions(
            base_options=BaseOptions(model_asset_path=model_path),
            min_detection_confidence=0.5,
        )
        _mp_detector = mp_vision.FaceDetector.create_from_options(options)
        print("MediaPipe FaceDetector ready.")
    except Exception as e:
        print(f"MediaPipe unavailable ({e}) — center-crop fallback will be used.")
        _mp_detector = None
    return _mp_detector


# ── Frame helpers ─────────────────────────────────────────────────────────────

def _center_crop(frame_rgb: np.ndarray, size: int) -> np.ndarray:
    H, W = frame_rgb.shape[:2]
    s    = min(H, W)
    top  = (H - s) // 2
    left = (W - s) // 2
    return cv2.resize(frame_rgb[top: top + s, left: left + s], (size, size))


def _mediapipe_crop(
    frame_rgb: np.ndarray,
    detector,
    margin: float,
    size: int,
) -> np.ndarray | None:
    try:
        import mediapipe as mp
        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame_rgb)
        result   = detector.detect(mp_image)
        if result.detections:
            H, W = frame_rgb.shape[:2]
            best = max(result.detections,
                       key=lambda d: d.bounding_box.width * d.bounding_box.height)
            bb   = best.bounding_box
            bw, bh = bb.width, bb.height
            mg   = int(max(bw, bh) * margin)
            x1   = max(0, bb.origin_x - mg)
            y1   = max(0, bb.origin_y - mg)
            x2   = min(W, bb.origin_x + bw + mg)
            y2   = min(H, bb.origin_y + bh + mg)
            crop = frame_rgb[y1:y2, x1:x2]
            if crop.size > 0:
                return cv2.resize(crop, (size, size))
    except Exception:
        pass
    return None


def extract_frames_from_bytes(video_bytes: bytes, num_frames: int = NUM_FRAMES) -> list:
    tmp_path = None
    try:
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f:
            f.write(video_bytes)
            tmp_path = f.name

        cap   = cv2.VideoCapture(tmp_path)
        total = max(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), 1)
        indices = np.linspace(0, total - 1, num_frames, dtype=int)
        frames  = []
        for idx in indices:
            cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
            ret, frame = cap.read()
            if ret and frame is not None:
                frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        cap.release()
        return frames
    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.unlink(tmp_path)


def process_frames(frames: list, detector) -> list:
    """Crop each frame to face (or center) and resize to FRAME_SIZE."""
    processed = []
    for f in frames:
        crop = None
        if detector is not None:
            crop = _mediapipe_crop(f, detector, FACE_MARGIN, FRAME_SIZE)
        if crop is None:
            crop = _center_crop(f, FRAME_SIZE)
        processed.append(crop)

    # Pad if we got fewer frames than expected (e.g. short video)
    while len(processed) < NUM_FRAMES:
        processed.append(processed[-1] if processed else
                         np.zeros((FRAME_SIZE, FRAME_SIZE, 3), dtype=np.uint8))
    return processed[:NUM_FRAMES]


# ── Main inference function ───────────────────────────────────────────────────

def get_visual_embedding(
    video_bytes: bytes,
    model,
    device: torch.device,
    face_detector=None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Run the full visual pipeline on raw video bytes.

    Returns
    -------
    preds : (5,)  float32 OCEAN predictions  [0, 1]
    emb   : (512,) float32 visual embedding
    """
    frames = extract_frames_from_bytes(video_bytes, NUM_FRAMES)
    if not frames:
        raise ValueError("Could not extract any frames from the video.")

    crops  = process_frames(frames, face_detector)
    tensor = torch.stack([FRAME_TFMS(c) for c in crops])   # (T, 3, 224, 224)
    tensor = tensor.unsqueeze(0).to(device)                 # (1, T, 3, 224, 224)

    model.eval()
    with torch.no_grad():
        preds, emb = model(tensor)

    return preds.squeeze(0).float().cpu(), emb.squeeze(0).float().cpu()