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)