File size: 4,418 Bytes
178f61f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | import os
import cv2
import numpy as np
import h5py
import scipy.io
from tqdm import tqdm
import sys
from pathlib import Path
# Add project root to path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
from src.utils.preprocess import GazePreprocessor
class RobustGazePreprocessor(GazePreprocessor):
def get_landmarks_robust(self, frame):
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
landmarks = self.get_landmarks(frame)
if landmarks: return landmarks
return None
def create_unseen_split():
mat_path = 'data/raw/metadata.mat'
img_root = 'data/raw/imgs'
output_path = 'data/processed/gaze360_unseen_verification.h5'
# Target IDs in Recording 6 that were NOT in test_B
# rec_006 IDs: [0, 1, 17, 25, 49, 60, 61, 62]
# test_B had: [60, 62]
UNSEEN_IDS = [0, 1, 17, 25, 49, 61]
REC_IDX = 6
patch_size = 16
print(f"Loading metadata and selecting unseen IDs: {UNSEEN_IDS}")
mat = scipy.io.loadmat(mat_path)
rec_indices = mat['recording'].flatten()
person_ids = mat['person_identity'].flatten()
mask = (rec_indices == REC_IDX) & np.isin(person_ids, UNSEEN_IDS)
indices = np.where(mask)[0]
rec_name = mat['recordings'][0, REC_IDX][0]
preprocessor = RobustGazePreprocessor()
os.makedirs('data/processed', exist_ok=True)
with h5py.File(output_path, 'w') as h5f:
lp_ds = h5f.create_dataset('left_patches', (0, 4, patch_size, patch_size), maxshape=(None, 4, patch_size, patch_size), dtype='uint8', compression='gzip')
rp_ds = h5f.create_dataset('right_patches', (0, 4, patch_size, patch_size), maxshape=(None, 4, patch_size, patch_size), dtype='uint8', compression='gzip')
gaze_ds = h5f.create_dataset('gaze', (0, 2), maxshape=(None, 2), dtype='float32')
lm_ds = h5f.create_dataset('landmarks', (0, 478, 2), maxshape=(None, 478, 2), dtype='float32')
count = 0
for i in tqdm(indices, desc="Preprocessing Unseen Subjects"):
p_id = person_ids[i]
frame_num = mat['frame'][0, i]
gaze_3d = mat['gaze_dir'][i]
img_path = os.path.join(img_root, rec_name, 'head', f'{p_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 = preprocessor.get_landmarks_robust(frame)
if landmarks is None: continue
try:
# Normalize and extract patches
left_eye_img, left_angle = preprocessor.normalize_eye(frame, landmarks, 'left', method='new')
left_patches = preprocessor.extract_patches(left_eye_img, patch_size=patch_size)
right_eye_img, right_angle = preprocessor.normalize_eye(frame, landmarks, 'right', method='new')
right_patches = preprocessor.extract_patches(right_eye_img, patch_size=patch_size)
# Align Gaze
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 relative to face center
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
lp_ds.resize((count + 1, 4, patch_size, patch_size))
lp_ds[count] = left_patches
rp_ds.resize((count + 1, 4, patch_size, patch_size))
rp_ds[count] = right_patches
gaze_ds.resize((count + 1, 2))
gaze_ds[count] = gaze_rad
lm_ds.resize((count + 1, 478, 2))
lm_ds[count] = landmarks_centered
count += 1
except:
continue
print(f"\nDone! Created {output_path} with {count} samples.")
if __name__ == "__main__":
create_unseen_split()
|