File size: 2,647 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 | import torch
import numpy as np
import h5py
import os
import sys
from pathlib import Path
from tqdm import tqdm
# Add project root to path
project_root = str(Path(__file__).parent.parent.parent)
if project_root not in sys.path:
sys.path.append(project_root)
from src.models.student import LIPEV2Student
def check_alignment(model_path, h5_path):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = LIPEV2Student().to(device)
state_dict = torch.load(model_path, map_location=device)
model.load_state_dict(state_dict)
model.eval()
# Metrics for different combinations
# (P_sign, Y_sign)
combinations = [
(1, 1), # Normal
(1, -1), # Inverted Yaw
(-1, 1), # Inverted Pitch
(-1, -1) # Both Inverted
]
errors = {c: 0.0 for c in combinations}
count = 0
with h5py.File(h5_path, 'r') as f:
lp = f['left_patches'][:]
rp = f['right_patches'][:]
lm = f['landmarks'][:]
g_gt = f['gaze'][:]
num_samples = min(500, lp.shape[0]) # Use subset for speed
with torch.no_grad():
for i in range(num_samples):
p_l, y_l = model(torch.from_numpy(lp[i]).float().unsqueeze(0).to(device),
torch.from_numpy(lm[i]).float().view(1, -1).to(device), state='A')
p_r, y_r = model(torch.from_numpy(rp[i]).float().unsqueeze(0).to(device),
torch.from_numpy(lm[i]).float().view(1, -1).to(device), state='A')
def l2d(p, y):
idx = torch.arange(90).float().to(device)
pp, yp = torch.softmax(p, 1), torch.softmax(y, 1)
return (torch.sum(pp*idx,1)*2-90), (torch.sum(yp*idx,1)*2-90)
pl, yl = l2d(p_l, y_l)
pr, yr = l2d(p_r, y_r)
pf, yf = (pl+pr)/2, (yl+yr)/2
gt_d = torch.from_numpy(g_gt[i]).to(device) * (180.0/np.pi)
for ps, ys in combinations:
err = (torch.abs(ps*pf - gt_d[0]) + torch.abs(ys*yf - gt_d[1])).item()
errors[(ps, ys)] += err
count += 1
print(f"\n" + "="*40)
print(f"{'SIGN (Pitch, Yaw)':<20} | {'MAE (deg)':<10}")
print(f"-"*40)
for c, err in errors.items():
mae = err / (count * 2)
print(f"{str(c):<20} | {mae:.4f}")
print(f"="*40)
if __name__ == "__main__":
check_alignment('checkpoints/best_student_p04.pt', 'data/processed/gaze360_robust_v16.h5')
|