File size: 2,539 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 | import torch
import numpy as np
import h5py
import os
import sys
from pathlib import Path
# Add project root to path
sys.path.append(str(Path(__file__).parent.parent))
from src.models.student import LIPEV2StudentGold
def debug_gaze360(model_path, h5_path):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = LIPEV2StudentGold().to(device)
state_dict = torch.load(model_path, map_location=device)
model.load_state_dict(state_dict, strict=False)
model.eval()
with h5py.File(h5_path, 'r') as f:
left_patches = f['left_patches'][:]
right_patches = f['right_patches'][:]
landmarks = f['landmarks'][:]
gaze_gt = f['gaze'][:] # (pitch, yaw) in radians
target_indices = list(range(10)) + [1004]
print(f"{'Sample':<6} | {'P_Pred':<8} | {'P_GT':<8} | {'Y_Pred':<8} | {'Y_GT':<8} | {'Error':<8}")
print("-" * 65)
total_err = 0
with torch.no_grad():
for i in target_indices:
if i >= left_patches.shape[0]: continue
lp = torch.from_numpy(left_patches[i]).float().unsqueeze(0).to(device) / 255.0
rp = torch.from_numpy(right_patches[i]).float().unsqueeze(0).to(device) / 255.0
lm = torch.from_numpy(landmarks[i]).float().view(1, -1).to(device)
gt = gaze_gt[i] * (180.0 / np.pi)
out_l = model(lp, lm)
out_r = model(rp, lm)
def logits_to_deg(out):
idx = torch.arange(90).float().to(device)
p_deg = (torch.sum(torch.softmax(out[0], dim=1) * idx, dim=1) * 2 - 90)
y_deg = (torch.sum(torch.softmax(out[1], dim=1) * idx, dim=1) * 2 - 90)
return p_deg.item(), y_deg.item()
p_l, y_l = logits_to_deg(out_l)
p_r, y_r = logits_to_deg(out_r)
p_pred = (p_l + p_r) / 2
y_pred = (y_l + y_r) / 2
p_gt, y_gt = gt[0], gt[1]
err = (abs(p_pred - p_gt) + abs(y_pred - y_gt)) / 2
total_err += err
print(f"{i:<6} | {p_pred:8.2f} | {p_gt:8.2f} | {y_pred:8.2f} | {y_gt:8.2f} | {err:8.2f}")
print("-" * 65)
print(f"Average of target samples: {total_err/len(target_indices):.4f}")
if __name__ == "__main__":
debug_gaze360('checkpoints/gold_swa/best_gold_p00.pt', 'data/processed/gaze360_robust_v16.h5')
|