File size: 4,284 Bytes
a10ba7f | 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 | import gradio as gr
import cv2
import torch
import numpy as np
from PIL import Image
import os
import sys
from pathlib import Path
# Add src to path
sys.path.append(str(Path(__file__).parent))
from src.models.student import LIPEV2StudentGaze360Gold
from src.utils.preprocess import GazePreprocessor
# 1. Model Initialization
device = "cuda" if torch.cuda.is_available() else "cpu"
model = LIPEV2StudentGaze360Gold()
# Load the best available checkpoint
checkpoint_path = "checkpoints/swa_gold_p11.pt"
if not os.path.exists(checkpoint_path):
# Fallback search
pts = list(Path("checkpoints").rglob("*.pt"))
if pts:
checkpoint_path = str(pts[0])
if os.path.exists(checkpoint_path):
checkpoint = torch.load(checkpoint_path, map_location=device)
model.load_state_dict(checkpoint)
print(f"Loaded {checkpoint_path}")
model.to(device)
model.eval()
# 2. Preprocessor Initialization
# Note: face_landmarker.task must be in the root directory for HF Spaces
preprocessor = GazePreprocessor(model_path='face_landmarker.task')
def draw_gaze(frame, pitch, yaw, landmarks):
h, w, _ = frame.shape
# Simple visualization logic
# Calculate eye centers (using indices from preprocessor)
left_eye_center = np.mean([[landmarks[i].x * w, landmarks[i].y * h] for i in [362, 263]], axis=0).astype(int)
right_eye_center = np.mean([[landmarks[i].x * w, landmarks[i].y * h] for i in [33, 133]], axis=0).astype(int)
p_rad = np.radians(pitch)
y_rad = np.radians(yaw)
dx = -np.sin(y_rad) * np.cos(p_rad)
dy = -np.sin(p_rad)
length = 80
for center in [left_eye_center, right_eye_center]:
end_point = (int(center[0] + dx * length), int(center[1] + dy * length))
cv2.arrowedLine(frame, tuple(center), end_point, (0, 255, 0), 2, tipLength=0.3)
return frame
def predict(img):
if img is None:
return None, "Please upload an image."
# Convert PIL to BGR
frame = np.array(img)
frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
# Get landmarks
landmarks = preprocessor.get_landmarks(frame)
if not landmarks:
return img, "No face detected."
# Prepare landmarks (zero-centered)
landmarks_arr = np.array([[lm.x, lm.y] for lm in landmarks])
left_c = np.mean([[landmarks[idx].x, landmarks[idx].y] for idx in [362, 263]], axis=0)
right_c = np.mean([[landmarks[idx].x, landmarks[idx].y] for idx in [33, 133]], axis=0)
face_center = (left_c + right_c) / 2
landmarks_centered = (landmarks_arr - face_center).flatten()
landmarks_tensor = torch.from_numpy(landmarks_centered).float().unsqueeze(0).to(device)
# Prepare patches (using left eye for inference)
left_eye_img, _ = preprocessor.normalize_eye(frame, landmarks, 'left')
patches = preprocessor.extract_patches(left_eye_img, patch_size=16)
patches_tensor = torch.from_numpy(patches).float().unsqueeze(0).to(device) / 255.0
# Inference
with torch.no_grad():
p_logits, y_logits, _ = model(patches_tensor, landmarks_tensor)
idx = torch.arange(90).float().to(device)
pitch = (torch.sum(torch.softmax(p_logits, dim=1) * idx, dim=1) * 2 - 90).item()
yaw = (torch.sum(torch.softmax(y_logits, dim=1) * idx, dim=1) * 2 - 90).item()
# Draw
res_frame = draw_gaze(frame.copy(), pitch, yaw, landmarks)
res_img = cv2.cvtColor(res_frame, cv2.COLOR_BGR2RGB)
return Image.fromarray(res_img), f"Pitch: {pitch:.2f}°, Yaw: {yaw:.2f}°"
# 3. Gradio Interface
description = """
### LIPE V2: Ultra-Lightweight Gaze Estimation
LIPE (Landmark-guided Image Patch Embedder) is a high-efficiency gaze estimation framework optimized for commodity CPUs.
- **Model Size:** 0.61M Parameters
- **Complexity:** 21.44 MFLOPs
- **Accuracy:** 4.73° (MPII), 7.33° (Gaze360 Zero-shot)
- **Speed:** 30 FPS on Single-threaded Mobile CPU
"""
iface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs=[gr.Image(label="Gaze Visualization"), gr.Textbox(label="Result")],
title="LIPE V2 Live Demo",
description=description,
examples=[["data/verification/sample_face.jpg"]] if os.path.exists("data/verification/sample_face.jpg") else None
)
if __name__ == "__main__":
iface.launch()
|