src/data/generate_teacher_labels.py

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)