src/data/preprocess_gaze360_robust.py

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
    )