src/data/extract_8x8_ablation.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

class AblationPreprocessor8x8(GazePreprocessor):
    def __init__(self, model_path='face_landmarker.task'):
        super().__init__(model_path)
        # Adjusted CLAHE for 8x8: smaller tile grid (2x2)
        self.clahe_8x8 = cv2.createCLAHE(clipLimit=1.2, tileGridSize=(2, 2))

    def normalize_eye_8x8_direct(self, frame, landmarks, eye_side='left'):
        """
        DIRECT 8x8 EXTRACTION:
        Warp directly to 32x16 eye ROI (exactly half of the 64x32 baseline).
        """
        h, w, _ = frame.shape
        indices = self.LEFT_CORNERS if eye_side == 'left' else self.RIGHT_CORNERS
            
        p1 = np.array([landmarks[indices[0]].x * w, landmarks[indices[0]].y * h])
        p2 = np.array([landmarks[indices[1]].x * w, landmarks[indices[1]].y * h])
        
        center = (p1 + p2) / 2
        dx, dy = p2 - p1
        angle = np.degrees(np.arctan2(dy, dx))
        
        dist = np.linalg.norm(p2 - p1)
        # Target size for 8x8 patches (K=4) is 32x16 for the whole eye
        target_size = (32, 16)
        scale = (target_size[0] * 0.7) / (dist + 1e-6)
        
        M = cv2.getRotationMatrix2D(tuple(center), angle, scale)
        M[0, 2] += (target_size[0] / 2) - center[0]
        M[1, 2] += (target_size[1] / 2) - center[1]
        
        # Warp directly to 32x16 using high-quality CUBIC interpolation
        normalized = cv2.warpAffine(frame, M, target_size, flags=cv2.INTER_CUBIC)
        normalized = cv2.cvtColor(normalized, cv2.COLOR_BGR2GRAY)

        # Step 2: Median Blur (same as baseline)
        normalized = cv2.medianBlur(normalized, 3)

        # Step 3: Adjusted CLAHE for small resolution
        normalized = self.clahe_8x8.apply(normalized)
            
        return normalized, angle

    def extract_patches_8x8_direct(self, eye_img):
        """
        Split 32x16 eye ROI into 4 quadrants (16x8) and resize to 8x8.
        Maintains the same 2:1 width-squish ratio as the 16x16 baseline.
        """
        h, w = eye_img.shape
        patches = []
        step_w, step_h = w // 2, h // 2
        
        for i in range(2):
            for j in range(2):
                roi = eye_img[i*step_h : (i+1)*step_h, j*step_w : (j+1)*step_w]
                # Resize 16x8 -> 8x8 using INTER_CUBIC (Fair comparison)
                patch = cv2.resize(roi, (8, 8), interpolation=cv2.INTER_CUBIC)
                patches.append(patch)
        return np.array(patches)

def parse_annotation(line):
    parts = line.split()
    if len(parts) < 41: return None
    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_8x8(p_id, data_root, output_dir, preprocessor):
    p_path = os.path.join(data_root, 'Data', 'Original', p_id)
    output_path = os.path.join(output_dir, f'{p_id}_8x8_ablation.h5')
    
    if not os.path.exists(p_path): return

    with h5py.File(output_path, 'w') as h5f:
        lp_ds = h5f.create_dataset('left_patches', (0, 4, 8, 8), maxshape=(None, 4, 8, 8), dtype='uint8', compression='gzip')
        rp_ds = h5f.create_dataset('right_patches', (0, 4, 8, 8), maxshape=(None, 4, 8, 8), dtype='uint8', compression='gzip')
        lg_ds = h5f.create_dataset('left_gaze', (0, 2), maxshape=(None, 2), dtype='float32')
        rg_ds = h5f.create_dataset('right_gaze', (0, 2), maxshape=(None, 2), dtype='float32')
        lm_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"Ablation 8x8: {p_id}"):
            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):
                ann = parse_annotation(line)
                if ann is None: continue
                img_path = os.path.join(day_path, f"{i+1:04d}.jpg")
                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
                
                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
                
                # Left Eye (Direct 8x8)
                le_img, le_angle = preprocessor.normalize_eye_8x8_direct(frame, landmarks, 'left')
                le_patches = preprocessor.extract_patches_8x8_direct(le_img)
                g_left = ann['target'] - ann['left_eye']
                g_left /= np.linalg.norm(g_left)
                gaze_left_rad = preprocessor.gaze_3d_to_mag(preprocessor.rotate_gaze(g_left, le_angle))
                
                # Right Eye (Direct 8x8)
                re_img, re_angle = preprocessor.normalize_eye_8x8_direct(frame, landmarks, 'right')
                re_patches = preprocessor.extract_patches_8x8_direct(re_img)
                g_right = ann['target'] - ann['right_eye']
                g_right /= np.linalg.norm(g_right)
                gaze_right_rad = preprocessor.gaze_3d_to_mag(preprocessor.rotate_gaze(g_right, re_angle))
                
                for ds, data in zip([lp_ds, rp_ds, lg_ds, rg_ds, lm_ds],
                                   [le_patches, re_patches, gaze_left_rad, gaze_right_rad, landmarks_centered]):
                    ds.resize((sample_idx + 1, *ds.shape[1:]))
                    ds[sample_idx] = data
                sample_idx += 1
    print(f"Finished {p_id}, samples: {sample_idx}")

if __name__ == '__main__':
    preprocessor = AblationPreprocessor8x8()
    participants = [f'p{i:02d}' for i in range(15)]
    for p_id in participants:
        process_participant_8x8(p_id, 'data/MPIIGaze/MPIIGaze/MPIIGaze', 'data/processed', preprocessor)