src/data/preprocess_mpii.py

import os
import cv2
import numpy as np
import h5py
from tqdm import tqdm
import sys
from pathlib import Path

# Add src to path
sys.path.append(str(Path(__file__).parent.parent.parent))
from src.utils.preprocess import GazePreprocessor

def parse_annotation(line):
    """
    Parse a line from MPIIGaze Data/Original/pXX/dayYY/annotation.txt
    """
    parts = line.split()
    if len(parts) < 41:
        return None
    
    # Corrected indices for Target and Eye Centers
    target_ccs = np.array([float(parts[26]), float(parts[27]), float(parts[28])])
    left_eye_ccs = np.array([float(parts[32]), float(parts[33]), float(parts[34])])
    right_eye_ccs = np.array([float(parts[35]), float(parts[36]), float(parts[37])])
    
    return {
        'target': target_ccs,
        'left_eye': left_eye_ccs,
        'right_eye': right_eye_ccs
    }

def process_participant(p_id, data_root, output_dir, preprocessor, patch_size=8, limit=None, method='new', suffix_extra=""):
    p_path = os.path.join(data_root, 'Data', 'Original', p_id)
    
    suffix = f"_v{patch_size}" if patch_size != 8 else ""
    output_path = os.path.join(output_dir, f'{p_id}{suffix}{suffix_extra}.h5')
    
    if not os.path.exists(p_path):
        print(f"Path not found: {p_path}")
        return

    # Prepare HDF5
    with h5py.File(output_path, 'w') as h5f:
        # We'll use resizable datasets
        left_patches_ds = h5f.create_dataset('left_patches', (0, 4, patch_size, patch_size), maxshape=(None, 4, patch_size, patch_size), dtype='uint8', compression='gzip')
        right_patches_ds = h5f.create_dataset('right_patches', (0, 4, patch_size, patch_size), maxshape=(None, 4, patch_size, patch_size), dtype='uint8', compression='gzip')
        left_gaze_ds = h5f.create_dataset('left_gaze', (0, 2), maxshape=(None, 2), dtype='float32')
        right_gaze_ds = h5f.create_dataset('right_gaze', (0, 2), maxshape=(None, 2), dtype='float32')
        landmarks_ds = h5f.create_dataset('landmarks', (0, 478, 2), maxshape=(None, 478, 2), dtype='float32')
        
        sample_idx = 0
        
        days = sorted([d for d in os.listdir(p_path) if d.startswith('day')])
        for day in tqdm(days, desc=f"Processing {p_id} ({method})"):
            day_path = os.path.join(p_path, 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 limit and sample_idx >= limit:
                    break
                    
                ann = parse_annotation(line)
                if ann is None:
                    continue
                
                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
                
                # Convert landmarks to numpy array (N, 2)
                landmarks_arr = np.array([[lm.x, lm.y] for lm in landmarks])
                
                # Landmark Zero-Centering
                left_c = np.mean([[landmarks[idx].x, landmarks[idx].y] for idx in preprocessor.LEFT_CORNERS], axis=0)
                right_c = np.mean([[landmarks[idx].x, landmarks[idx].y] for idx in preprocessor.RIGHT_CORNERS], axis=0)
                face_center = (left_c + right_c) / 2
                landmarks_centered = landmarks_arr - face_center
                
                # Process Left Eye
                left_eye_img, left_angle = preprocessor.normalize_eye(frame, landmarks, 'left', method=method)
                left_patches = preprocessor.extract_patches(left_eye_img, patch_size=patch_size)
                
                # Gaze for Left Eye
                g_left = ann['target'] - ann['left_eye']
                g_left /= np.linalg.norm(g_left)
                g_left_rot = preprocessor.rotate_gaze(g_left, left_angle)
                gaze_left_rad = preprocessor.gaze_3d_to_mag(g_left_rot)
                
                # Process Right Eye
                right_eye_img, right_angle = preprocessor.normalize_eye(frame, landmarks, 'right', method=method)
                right_patches = preprocessor.extract_patches(right_eye_img, patch_size=patch_size)
                
                # Gaze for Right Eye
                g_right = ann['target'] - ann['right_eye']
                g_right /= np.linalg.norm(g_right)
                g_right_rot = preprocessor.rotate_gaze(g_right, right_angle)
                gaze_right_rad = preprocessor.gaze_3d_to_mag(g_right_rot)
                
                # Append to HDF5
                for ds, data in zip([left_patches_ds, right_patches_ds, left_gaze_ds, right_gaze_ds, landmarks_ds],
                                   [left_patches, right_patches, gaze_left_rad, gaze_right_rad, landmarks_centered]):
                    ds.resize((sample_idx + 1, *ds.shape[1:]))
                    ds[sample_idx] = data
                
                sample_idx += 1
            if limit and sample_idx >= limit:
                break

    print(f"Finished {p_id}, total samples: {sample_idx}")

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_root', type=str, default='data/MPIIGaze/MPIIGaze/MPIIGaze')
    parser.add_argument('--output_dir', type=str, default='data/processed')
    parser.add_argument('--participants', type=str, default='all', help='all or p00,p01...')
    parser.add_argument('--patch_size', type=int, default=8)
    parser.add_argument('--method', type=str, default='new', choices=['old', 'new'])
    parser.add_argument('--suffix_extra', type=str, default='_new')
    args = parser.parse_args()
    
    os.makedirs(args.output_dir, exist_ok=True)
    
    preprocessor = GazePreprocessor()
    
    if args.participants == 'all':
        participants = [f'p{i:02d}' for i in range(15)]
    else:
        participants = args.participants.split(',')
        
    for p_id in participants:
        process_participant(p_id, args.data_root, args.output_dir, preprocessor, patch_size=args.patch_size, method=args.method, suffix_extra=args.suffix_extra)