File size: 5,934 Bytes
396c01a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | 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()
|