| import torch |
| import torch.nn as nn |
| import numpy as np |
| import h5py |
| import os |
| import sys |
| from pathlib import Path |
| from tqdm import tqdm |
|
|
| |
| class MiniConvEmbedder(nn.Module): |
| def __init__(self): |
| super(MiniConvEmbedder, self).__init__() |
| self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=0) |
| self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=0) |
| self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=0) |
| self.relu = nn.ReLU(inplace=True) |
| self.gap = nn.AdaptiveAvgPool2d(1) |
| def forward(self, x): |
| x = self.relu(self.conv1(x)) |
| x = self.relu(self.conv2(x)) |
| x = self.relu(self.conv3(x)) |
| x = self.gap(x) |
| return torch.flatten(x, 1) |
|
|
| class LIPEV2StudentBaseline(nn.Module): |
| def __init__(self): |
| super(LIPEV2StudentBaseline, self).__init__() |
| self.appearance_net = MiniConvEmbedder() |
| self.geo_mlp = nn.Sequential( |
| nn.Linear(956, 256), |
| nn.LayerNorm(256), |
| nn.ReLU(inplace=True), |
| nn.Dropout(0.05), |
| nn.Linear(256, 256), |
| nn.ReLU(inplace=True) |
| ) |
| self.pitch_head = nn.Sequential(nn.Linear(256, 64), nn.ReLU(inplace=True), nn.Linear(64, 90)) |
| self.yaw_head = nn.Sequential(nn.Linear(256, 64), nn.ReLU(inplace=True), nn.Linear(64, 90)) |
|
|
| def forward(self, patches=None, landmarks=None, state='A'): |
| geo_feat = self.geo_mlp(landmarks) |
| if state == 'A' and patches is not None: |
| batch_size = patches.shape[0] |
| patches = patches.view(-1, 1, patches.shape[2], patches.shape[3]) |
| app_tokens = self.appearance_net(patches) |
| app_feat = app_tokens.view(batch_size, -1) |
| combined = app_feat + geo_feat |
| else: |
| combined = geo_feat |
| return self.pitch_head(combined), self.yaw_head(combined) |
|
|
| def evaluate_baseline_on_gaze360(model_path, h5_path): |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| print(f"Evaluating BASELINE model: {model_path}") |
| |
| model = LIPEV2StudentBaseline().to(device) |
| state_dict = torch.load(model_path, map_location=device) |
| model.load_state_dict(state_dict) |
| model.eval() |
|
|
| results = { |
| 'all': {'error': 0.0, 'count': 0}, |
| 'frontal_45': {'error': 0.0, 'count': 0} |
| } |
|
|
| with h5py.File(h5_path, 'r') as f: |
| lp, rp, lm, g_gt = f['left_patches'][:], f['right_patches'][:], f['landmarks'][:], f['gaze'][:] |
| with torch.no_grad(): |
| for i in tqdm(range(len(lp)), desc="Testing Baseline"): |
| 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) |
| yaw_gt_deg = gt_d[1].item() |
| |
| error = (torch.abs(pf-gt_d[0]) + torch.abs(yf-gt_d[1])).item() |
| |
| results['all']['error'] += error |
| results['all']['count'] += 1 |
| if abs(yaw_gt_deg) <= 45.0: |
| results['frontal_45']['error'] += error |
| results['frontal_45']['count'] += 1 |
|
|
| print(f"\n" + "="*45) |
| print(f"{'SUBSET (BASELINE)':<20} | {'SAMPLES':<10} | {'MAE (deg)':<10}") |
| print(f"-"*45) |
| for key, data in results.items(): |
| if data['count'] > 0: |
| mae = data['error'] / (data['count'] * 2) |
| name = "All Cases" if key == 'all' else "Frontal +/- 45" |
| print(f"{name:<20} | {data['count']:<10} | {mae:.4f}") |
| print(f"="*45) |
|
|
| if __name__ == "__main__": |
| |
| evaluate_baseline_on_gaze360('checkpoints/baseline_v16/best_student_p08.pt', 'data/processed/gaze360_robust_v16_test_B.h5') |
|
|
