import os import cv2 import numpy as np import h5py import scipy.io from tqdm import tqdm import sys from pathlib import Path from dataclasses import dataclass # Add src to path sys.path.append(str(Path(__file__).parent.parent.parent)) from src.utils.preprocess import GazePreprocessor @dataclass class MockLandmark: x: float y: float z: float = 0.0 class RobustGazePreprocessor(GazePreprocessor): def get_landmarks_robust(self, frame): """ Thử tìm mặt ở nhiều góc xoay khác nhau (0, 180, 90, 270) để đảm bảo không bỏ sót ảnh nào MediaPipe có thể đọc được. """ rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) h, w = frame.shape[:2] # 1. Thử góc thẳng landmarks = self.get_landmarks(frame) if landmarks: return landmarks, 0 # 2. Thử các góc xoay (ưu tiên 180 độ cho trường hợp lộn đầu) for angle in [180, 90, 270]: if angle == 180: rotated = cv2.rotate(rgb_frame, cv2.ROTATE_180) elif angle == 90: rotated = cv2.rotate(rgb_frame, cv2.ROTATE_90_CLOCKWISE) else: rotated = cv2.rotate(rgb_frame, cv2.ROTATE_90_COUNTERCLOCKWISE) # Gọi hàm detect gốc của MediaPipe import mediapipe as mp mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rotated) res = self.detector.detect(mp_image) if res.face_landmarks: # Nắn tọa độ landmark về khung hình gốc lms = res.face_landmarks[0] unrotated = [] for lm in lms: x, y = lm.x, lm.y if angle == 180: nx, ny = 1.0 - x, 1.0 - y elif angle == 90: nx, ny = y, 1.0 - x elif angle == 270: nx, ny = 1.0 - y, x unrotated.append(MockLandmark(x=nx, y=ny, z=lm.z)) return unrotated, angle return None, 0 def run_robust_processing(output_prefix='data/processed/gaze360_v16', method='new', train_ids=None, test_ids=None): mat_path = 'data/raw/metadata.mat' img_root = 'data/raw/imgs' rec_idx = 6 patch_size = 16 print(f"Loading metadata... [Method: {method}]") mat = scipy.io.loadmat(mat_path) indices = np.where(mat['recording'].flatten() == rec_idx)[0] rec_name = mat['recordings'][0, rec_idx][0] preprocessor = RobustGazePreprocessor() os.makedirs(os.path.dirname(output_prefix), exist_ok=True) # Train and Test H5 files train_path = f"{output_prefix}_train_A.h5" test_path = f"{output_prefix}_test_B.h5" with h5py.File(train_path, 'w') as h5_train, h5py.File(test_path, 'w') as h5_test: datasets = {} for mode, h5f in zip(['train', 'test'], [h5_train, h5_test]): datasets[mode] = { 'lp': h5f.create_dataset('left_patches', (0, 4, patch_size, patch_size), maxshape=(None, 4, patch_size, patch_size), dtype='uint8', compression='gzip'), 'rp': h5f.create_dataset('right_patches', (0, 4, patch_size, patch_size), maxshape=(None, 4, patch_size, patch_size), dtype='uint8', compression='gzip'), 'gaze': h5f.create_dataset('gaze', (0, 2), maxshape=(None, 2), dtype='float32'), 'lm': h5f.create_dataset('landmarks', (0, 478, 2), maxshape=(None, 478, 2), dtype='float32'), 'idx': 0 } for i in tqdm(indices, desc=f"Robust Splitting ({method})"): person_id = mat['person_identity'][0, i] # Determine which set this subject belongs to if train_ids is not None and person_id in train_ids: mode = 'train' elif test_ids is not None and person_id in test_ids: mode = 'test' else: continue # Skip subjects not in our targeted split frame_num = mat['frame'][0, i] gaze_3d = mat['gaze_dir'][i] img_path = os.path.join(img_root, rec_name, 'head', f'{person_id:06d}', f'{frame_num:06d}.jpg') if not os.path.exists(img_path): continue frame = cv2.imread(img_path) if frame is None: continue landmarks, det_angle = preprocessor.get_landmarks_robust(frame) if landmarks is None: continue try: 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) 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) avg_angle = (left_angle + right_angle) / 2 gaze_rot = preprocessor.rotate_gaze(gaze_3d, avg_angle) gaze_rad = preprocessor.gaze_3d_to_mag(gaze_rot) # Landmarks zero-centering lms_arr = np.array([[lm.x, lm.y] for lm in landmarks]) 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 = lms_arr - face_center # Save to specific H5 ds = datasets[mode] cur_idx = ds['idx'] ds['lp'].resize((cur_idx + 1, 4, patch_size, patch_size)) ds['lp'][cur_idx] = left_patches ds['rp'].resize((cur_idx + 1, 4, patch_size, patch_size)) ds['rp'][cur_idx] = right_patches ds['gaze'].resize((cur_idx + 1, 2)) ds['gaze'][cur_idx] = gaze_rad ds['lm'].resize((cur_idx + 1, 478, 2)) ds['lm'][cur_idx] = landmarks_centered ds['idx'] += 1 except Exception as e: continue print(f"\nDone! Method '{method}'") print(f"Train A: {datasets['train']['idx']} samples") print(f"Test B: {datasets['test']['idx']} samples") if __name__ == "__main__": # Define Split (Based on IDs identified for Recording 6) # IDs: [0, 1, 17, 25, 49, 60, 61, 62] # Balanced split to give more data to Train A (ID 61 is heavy) TRAIN_IDS = [0, 1, 17, 25, 49, 61] TEST_IDS = [60, 62] # Generate NEW (Huy's 4-step) version with Splitting run_robust_processing( output_prefix='data/processed/gaze360_robust_v16', method='new', train_ids=TRAIN_IDS, test_ids=TEST_IDS )