import os import cv2 import numpy as np import h5py import torch from tqdm import tqdm import sys 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)): """ Crop the face from the frame using landmarks. """ h, w, _ = frame.shape # Extract landmark coordinates coords = np.array([[lm.x * w, lm.y * h] for lm in landmarks]) # Get bounding box min_x, min_y = np.min(coords, axis=0) max_x, max_y = np.max(coords, axis=0) # Add padding and make it square 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 # 50% padding 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) # Normalize for ResNet (ImageNet stats) 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 # HWC -> CHW face_img = np.transpose(face_img, (2, 0, 1)) return face_img def generate_labels(data_root, processed_dir, checkpoint_path, device='cpu'): # Load Teacher model = load_teacher_model(checkpoint_path, device=device) model.eval() preprocessor = GazePreprocessor() # Get all processed .h5 files h5_files = sorted([f for f in os.listdir(processed_dir) if f.endswith('.h5')]) for h5_name in h5_files: p_id = h5_name.split('.')[0] h5_path = os.path.join(processed_dir, h5_name) # Original data path p_data_root = os.path.join(data_root, 'Data', 'Original', p_id) if not os.path.exists(p_data_root): print(f"Original data for {p_id} not found at {p_data_root}. Skipping.") continue print(f"Generating teacher labels for {p_id}...") # We'll read the existing HDF5 to know which frames were successfully processed # Note: preprocess_mpii.py skips frames where landmarks are not found. # We need to match the exact same frames. # Open HDF5 in append mode with h5py.File(h5_path, 'a') as h5f: num_samples = h5f['landmarks'].shape[0] # Create or overwrite datasets for logits # L2CS outputs 90 bins for pitch and 90 for yaw if 'teacher_pitch_logits' in h5f: del h5f['teacher_pitch_logits'] if 'teacher_yaw_logits' in h5f: del h5f['teacher_yaw_logits'] pitch_logits_ds = h5f.create_dataset('teacher_pitch_logits', (num_samples, 90), dtype='float32') yaw_logits_ds = h5f.create_dataset('teacher_yaw_logits', (num_samples, 90), dtype='float32') # Re-run the same logic as preprocess_mpii.py to find the images sample_idx = 0 days = sorted([d for d in os.listdir(p_data_root) if d.startswith('day')]) pbar = tqdm(total=num_samples, desc=f"Processing {p_id}") for day in days: day_path = os.path.join(p_data_root, day) ann_file = os.path.join(day_path, 'annotation.txt') if not os.path.exists(ann_file): continue with open(ann_file, 'r') as f: lines = f.readlines() for i, line in enumerate(lines): if sample_idx >= num_samples: break img_name = f"{i+1:04d}.jpg" img_path = os.path.join(day_path, img_name) if not os.path.exists(img_path): continue frame = cv2.imread(img_path) if frame is None: continue landmarks = preprocessor.get_landmarks(frame) if landmarks is None: continue # Get face crop for Teacher face_input = get_face_crop(frame, landmarks) if face_input is None: # This shouldn't really happen if landmarks were found, but just in case pitch_logits_ds[sample_idx] = np.zeros(90) yaw_logits_ds[sample_idx] = np.zeros(90) else: # Run Teacher input_tensor = torch.from_numpy(face_input).unsqueeze(0).to(device) with torch.no_grad(): p_logits, y_logits = model(input_tensor) pitch_logits_ds[sample_idx] = p_logits.cpu().numpy() yaw_logits_ds[sample_idx] = y_logits.cpu().numpy() sample_idx += 1 pbar.update(1) if sample_idx >= num_samples: break pbar.close() if __name__ == '__main__': import argparse parser = argparse.ArgumentParser() parser.add_argument('--data_root', type=str, default='data/MPIIGaze/MPIIGaze/MPIIGaze') parser.add_argument('--processed_dir', type=str, default='data/processed') parser.add_argument('--checkpoint_path', type=str, default='checkpoints/resnet50.pt') parser.add_argument('--participant', type=str, default=None, help='Specific participant pXX') args = parser.parse_args() # Use GPU if available device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f"Using device: {device}") if args.participant: h5_files = [f"{args.participant}.h5"] # Filter files that exist h5_files = [f for f in h5_files if os.path.exists(os.path.join(args.processed_dir, f))] # Load Teacher model = load_teacher_model(args.checkpoint_path, device=device) model.eval() preprocessor = GazePreprocessor() for h5_name in h5_files: p_id = h5_name.split('_')[0].split('.')[0] # Extract pXX from pXX.h5 or pXX_v16.h5 h5_path = os.path.join(args.processed_dir, h5_name) p_data_root = os.path.join(args.data_root, 'Data', 'Original', p_id) with h5py.File(h5_path, 'a') as h5f: num_samples = h5f['landmarks'].shape[0] if 'teacher_pitch_logits' in h5f: del h5f['teacher_pitch_logits'] if 'teacher_yaw_logits' in h5f: del h5f['teacher_yaw_logits'] pitch_logits_ds = h5f.create_dataset('teacher_pitch_logits', (num_samples, 90), dtype='float32') yaw_logits_ds = h5f.create_dataset('teacher_yaw_logits', (num_samples, 90), dtype='float32') sample_idx = 0 days = sorted([d for d in os.listdir(p_data_root) if d.startswith('day')]) pbar = tqdm(total=num_samples, desc=f"Processing {p_id}") for day in days: day_path = os.path.join(p_data_root, day) ann_file = os.path.join(day_path, 'annotation.txt') if not os.path.exists(ann_file): continue with open(ann_file, 'r') as f: lines = f.readlines() for i, line in enumerate(lines): if sample_idx >= num_samples: break img_name = f"{i+1:04d}.jpg" img_path = os.path.join(day_path, img_name) if not os.path.exists(img_path): continue frame = cv2.imread(img_path) if frame is None: continue landmarks = preprocessor.get_landmarks(frame) if landmarks is None: continue face_input = get_face_crop(frame, landmarks) if face_input is not None: input_tensor = torch.from_numpy(face_input).unsqueeze(0).to(device) with torch.no_grad(): p_logits, y_logits = model(input_tensor) pitch_logits_ds[sample_idx] = p_logits.cpu().numpy() yaw_logits_ds[sample_idx] = y_logits.cpu().numpy() sample_idx += 1 pbar.update(1) if sample_idx >= num_samples: break pbar.close() else: generate_labels(args.data_root, args.processed_dir, args.checkpoint_path, device=device)