File size: 4,466 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 | import os
import cv2
import numpy as np
import h5py
import scipy.io
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 process_gaze360_subset(mat_path, img_root, output_path, rec_idx=6, patch_size=16):
print(f"Loading metadata from {mat_path}...")
mat = scipy.io.loadmat(mat_path)
recording_indices = mat['recording'].flatten()
mask = (recording_indices == rec_idx)
indices = np.where(mask)[0]
print(f"Found {len(indices)} potential samples in recording {rec_idx}")
preprocessor = GazePreprocessor()
# Get recording name
rec_name = mat['recordings'][0, rec_idx][0]
os.makedirs(os.path.dirname(output_path), exist_ok=True)
with h5py.File(output_path, 'w') as h5f:
# 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')
# We store person's gaze as both left and right for compatibility
gaze_ds = h5f.create_dataset('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
for i in tqdm(indices, desc="Processing Gaze360"):
person_id = mat['person_identity'][0, i]
frame_num = mat['frame'][0, i]
gaze_3d = mat['gaze_dir'][i] # (3,)
# Construct path
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 = preprocessor.get_landmarks(frame)
if landmarks is None:
continue
# 1. Landmarks (Centered)
landmarks_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 = landmarks_arr - face_center
# 2. Eye Patches
try:
# We use Left Eye for normalizing the 3D gaze (assuming they are close)
# In a more precise setup, we'd do it for each eye.
left_eye_img, left_angle = preprocessor.normalize_eye(frame, landmarks, 'left')
left_patches = preprocessor.extract_patches(left_eye_img, patch_size=patch_size)
right_eye_img, right_angle = preprocessor.normalize_eye(frame, landmarks, 'right')
right_patches = preprocessor.extract_patches(right_eye_img, patch_size=patch_size)
# Rotate Gaze (Using average angle or just left eye's angle)
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) # (pitch, yaw)
# Store
for ds, data in zip([left_patches_ds, right_patches_ds, gaze_ds, landmarks_ds],
[left_patches, right_patches, gaze_rad, landmarks_centered]):
ds.resize((sample_idx + 1, *ds.shape[1:]))
ds[sample_idx] = data
sample_idx += 1
except Exception as e:
# print(f"Error processing frame: {e}")
continue
print(f"\nFinished! Processed {sample_idx} samples. Saved to {output_path}")
if __name__ == "__main__":
process_gaze360_subset(
mat_path='data/raw/metadata.mat',
img_root='data/raw/imgs',
output_path='data/processed/gaze360_v16.h5',
rec_idx=6,
patch_size=16
)
|