File size: 3,457 Bytes
178f61f | 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 | import torch
import cv2
import numpy as np
import sys
import os
from pathlib import Path
# Add project root to path
sys.path.append(str(Path(__file__).parent.parent.parent))
from src.models.teacher import load_teacher_model
from src.utils.preprocess import GazePreprocessor
def get_face_crop(frame, landmarks, target_size=(224, 224)):
h, w, _ = frame.shape
coords = np.array([[lm.x * w, lm.y * h] for lm in landmarks])
min_x, min_y = np.min(coords, axis=0)
max_x, max_y = np.max(coords, axis=0)
width = max_x - min_x
height = max_y - min_y
center_x = (min_x + max_x) / 2
center_y = (min_y + max_y) / 2
size = max(width, height) * 1.5
x1 = int(max(0, center_x - size / 2))
y1 = int(max(0, center_y - size / 2))
x2 = int(min(w, center_x + size / 2))
y2 = int(min(h, center_y + size / 2))
face_img = frame[y1:y2, x1:x2]
if face_img.size == 0: return None
face_img = cv2.resize(face_img, target_size)
face_img = cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB)
face_img = face_img.astype(np.float32) / 255.0
mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
face_img = (face_img - mean) / std
face_img = np.transpose(face_img, (2, 0, 1))
return face_img, (x1, y1, x2, y2)
def draw_gaze(image, pitch, yaw, length=100, color=(0, 255, 0)):
# Simple 2D projection for visualization
# Pitch (up/down), Yaw (left/right)
# We assume gaze starts at center
h, w = image.shape[:2]
cx, cy = w // 2, h // 2
# Negative yaw because +x is right, but gaze x is usually -x
dx = -length * np.sin(np.radians(yaw))
dy = -length * np.sin(np.radians(pitch)) # -pitch because +y is down
cv2.line(image, (cx, cy), (int(cx + dx), int(cy + dy)), color, 3)
return image
def visualize_teacher():
checkpoint_path = 'checkpoints/resnet50.pt'
# Use absolute path for MediaPipe model
model_path = os.path.abspath('face_landmarker.task')
device = 'cpu'
model = load_teacher_model(checkpoint_path, backbone='resnet50', device=device)
preprocessor = GazePreprocessor(model_path=model_path)
# Try multiple images from p00/day01
img_dir = 'data/MPIIGaze/MPIIGaze/MPIIGaze/Data/Original/p00/day01'
for i in range(1, 10):
img_path = os.path.join(img_dir, f"{i:04d}.jpg")
if not os.path.exists(img_path): continue
print(f"Testing on {img_path}...")
frame = cv2.imread(img_path)
landmarks = preprocessor.get_landmarks(frame)
if landmarks is None:
print(f"No landmarks found for {img_path}.")
continue
face_input, (x1, y1, x2, y2) = get_face_crop(frame, landmarks)
input_tensor = torch.from_numpy(face_input).unsqueeze(0).to(device)
with torch.no_grad():
p_logits, y_logits = model(input_tensor)
p_deg, y_deg = model.get_angles(p_logits, y_logits)
print(f"Predicted: Pitch {p_deg.item():.2f}, Yaw {y_deg.item():.2f}")
# Visualize
face_vis = frame[y1:y2, x1:x2].copy()
face_vis = draw_gaze(face_vis, p_deg.item(), y_deg.item())
output_path = f'data/verification/teacher_test_{i}.png'
cv2.imwrite(output_path, face_vis)
print(f"Saved to {output_path}")
break
if __name__ == '__main__':
visualize_teacher()
|