Gaze-LIPE / src /data /preprocess_gaze360.py

Publish KD reproducibility investigation

178f61f 3 months ago

4.47 kB

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
    )

	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
	)