Upload src/data_generator.py with huggingface_hub

src/data_generator.py

@@ -0,0 +1,425 @@
+"""
+Synthetic Data Generator for Eye Gaze Training
+
+Generates realistic synthetic training data that simulates:
+1. Eye crops with various iris positions (gaze directions)
+2. Dark / low-light conditions
+3. Glasses overlays
+4. Lazy eye / strabismus (asymmetric eye gaze)
+5. Various skin tones, eye colors
+6. Sensor noise (CMOS simulation)
+7. Illumination perturbation (directional light gradients)
+
+Based on augmentation strategies from:
+- AGE framework (arxiv:2603.26945) - GlassesGAN, illumination perturbation, sensor noise
+- UnityEyes approach - synthetic eye rendering with parametric control
+"""
+
+import numpy as np
+import tensorflow as tf
+from PIL import Image, ImageDraw, ImageFilter
+import random
+import math
+
+
+class SyntheticGazeDataGenerator:
+    """
+    Generates synthetic eye + face images with known gaze labels.
+    
+    Each sample contains:
+      - left_eye: 64x64 RGB crop
+      - right_eye: 64x64 RGB crop
+      - face: 64x64 RGB crop
+      - gaze_x, gaze_y: normalized screen coordinates [0, 1]
+    """
+    
+    def __init__(self, img_size=64, seed=42):
+        self.img_size = img_size
+        self.rng = np.random.RandomState(seed)
+        
+        # Skin tone palette (RGB) - diverse range
+        self.skin_tones = [
+            (255, 224, 189), (255, 205, 148), (234, 192, 134),
+            (255, 173, 96),  (210, 153, 83),  (187, 131, 71),
+            (156, 102, 52),  (128, 80, 37),   (100, 64, 30),
+            (74, 46, 21),    (60, 38, 18),    (45, 30, 15),
+        ]
+        
+        # Eye colors (iris RGB)
+        self.eye_colors = [
+            (50, 30, 10),    # Dark brown
+            (100, 60, 20),   # Light brown
+            (40, 80, 40),    # Green
+            (30, 50, 100),   # Blue
+            (50, 50, 50),    # Grey
+            (80, 40, 10),    # Hazel
+            (20, 20, 20),    # Very dark (common in Asian/African eyes)
+        ]
+        
+        # Glasses frame colors
+        self.glasses_colors = [
+            (0, 0, 0),       # Black
+            (60, 40, 20),    # Brown
+            (100, 100, 100), # Silver/grey
+            (0, 0, 60),      # Dark blue
+            (80, 0, 0),      # Dark red
+        ]
+    
+    def _draw_eye(self, gaze_x, gaze_y, skin_tone, eye_color, eye_openness=1.0, 
+                  lazy_offset_x=0.0, lazy_offset_y=0.0):
+        """Draw a synthetic eye with iris at position determined by gaze."""
+        size = self.img_size
+        img = Image.new('RGB', (size, size), skin_tone)
+        draw = ImageDraw.Draw(img)
+        
+        cx, cy = size // 2, size // 2
+        
+        # Eye white (sclera) - elliptical shape
+        eye_w = int(size * 0.75)
+        eye_h = int(size * 0.35 * eye_openness)
+        sclera_bbox = [cx - eye_w//2, cy - eye_h//2, cx + eye_w//2, cy + eye_h//2]
+        draw.ellipse(sclera_bbox, fill=(240, 240, 240), outline=(180, 150, 130))
+        
+        # Iris position: map gaze (0-1) to iris displacement within eye
+        # Gaze (0,0) = top-left of screen, (1,1) = bottom-right
+        # When looking left on screen, iris moves left relative to eye
+        max_disp_x = eye_w * 0.25
+        max_disp_y = eye_h * 0.2
+        
+        iris_offset_x = (gaze_x - 0.5) * 2 * max_disp_x + lazy_offset_x * max_disp_x
+        iris_offset_y = (gaze_y - 0.5) * 2 * max_disp_y + lazy_offset_y * max_disp_y
+        
+        iris_cx = cx + iris_offset_x
+        iris_cy = cy + iris_offset_y
+        iris_r = int(size * 0.14)
+        
+        # Draw iris
+        draw.ellipse([iris_cx - iris_r, iris_cy - iris_r, 
+                      iris_cx + iris_r, iris_cy + iris_r], fill=eye_color)
+        
+        # Draw pupil (darker center)
+        pupil_r = iris_r // 2
+        draw.ellipse([iris_cx - pupil_r, iris_cy - pupil_r,
+                      iris_cx + pupil_r, iris_cy + pupil_r], fill=(5, 5, 5))
+        
+        # Specular highlight (light reflection)
+        spec_r = max(2, iris_r // 4)
+        spec_x = iris_cx - iris_r * 0.3
+        spec_y = iris_cy - iris_r * 0.3
+        draw.ellipse([spec_x - spec_r, spec_y - spec_r,
+                      spec_x + spec_r, spec_y + spec_r], fill=(255, 255, 255))
+        
+        # Upper eyelid
+        lid_pts_upper = []
+        for i in range(20):
+            t = i / 19.0
+            x = sclera_bbox[0] + t * eye_w
+            # Parabolic eyelid shape
+            y = cy - eye_h//2 - int(eye_h * 0.2 * math.sin(t * math.pi))
+            lid_pts_upper.append((x, y))
+        lid_pts_upper.extend([(sclera_bbox[2], 0), (sclera_bbox[0], 0)])
+        draw.polygon(lid_pts_upper, fill=skin_tone)
+        
+        # Lower eyelid
+        lid_pts_lower = []
+        for i in range(20):
+            t = i / 19.0
+            x = sclera_bbox[0] + t * eye_w
+            y = cy + eye_h//2 + int(eye_h * 0.15 * math.sin(t * math.pi))
+            lid_pts_lower.append((x, y))
+        lid_pts_lower.extend([(sclera_bbox[2], size), (sclera_bbox[0], size)])
+        draw.polygon(lid_pts_lower, fill=skin_tone)
+        
+        # Eyelashes (thin lines)
+        for i in range(0, eye_w, 4):
+            x = sclera_bbox[0] + i
+            t = i / eye_w
+            y_base = cy - eye_h//2 - int(eye_h * 0.2 * math.sin(t * math.pi))
+            draw.line([(x, y_base), (x + self.rng.randint(-2, 3), y_base - self.rng.randint(2, 6))], 
+                     fill=(20, 15, 10), width=1)
+        
+        # Add slight blur for realism
+        img = img.filter(ImageFilter.GaussianBlur(radius=0.5))
+        
+        return np.array(img, dtype=np.float32)
+    
+    def _draw_face(self, skin_tone):
+        """Draw a simplified face crop (head pose context)."""
+        size = self.img_size
+        img = Image.new('RGB', (size, size), skin_tone)
+        draw = ImageDraw.Draw(img)
+        
+        cx, cy = size // 2, size // 2
+        
+        # Face oval
+        face_w, face_h = int(size * 0.8), int(size * 0.9)
+        draw.ellipse([cx - face_w//2, cy - face_h//2, cx + face_w//2, cy + face_h//2], 
+                     fill=skin_tone)
+        
+        # Eyebrow regions (darker)
+        darker = tuple(max(0, c - 40) for c in skin_tone)
+        draw.arc([cx - face_w//3, cy - face_h//4, cx - face_w//10, cy - face_h//6],
+                 180, 360, fill=darker, width=2)
+        draw.arc([cx + face_w//10, cy - face_h//4, cx + face_w//3, cy - face_h//6],
+                 180, 360, fill=darker, width=2)
+        
+        # Nose hint
+        draw.line([(cx, cy - face_h//8), (cx, cy + face_h//8)], fill=darker, width=1)
+        
+        # Mouth
+        draw.arc([cx - face_w//6, cy + face_h//6, cx + face_w//6, cy + face_h//4],
+                 0, 180, fill=(180, 80, 80), width=2)
+        
+        img = img.filter(ImageFilter.GaussianBlur(radius=1))
+        return np.array(img, dtype=np.float32)
+    
+    def _add_glasses(self, eye_img, glasses_color):
+        """Overlay glasses frame on eye image."""
+        img = Image.fromarray(eye_img.astype(np.uint8))
+        draw = ImageDraw.Draw(img)
+        size = self.img_size
+        cx, cy = size // 2, size // 2
+        
+        # Frame outline (circular lens)
+        r = int(size * 0.35)
+        frame_width = self.rng.randint(2, 5)
+        draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=glasses_color, width=frame_width)
+        
+        # Temple arm hint
+        draw.line([(cx + r, cy), (size, cy - 2)], fill=glasses_color, width=frame_width)
+        
+        # Lens tint/reflection (subtle)
+        if self.rng.random() > 0.5:
+            overlay = Image.new('RGBA', (size, size), (0, 0, 0, 0))
+            overlay_draw = ImageDraw.Draw(overlay)
+            tint_alpha = self.rng.randint(10, 40)
+            overlay_draw.ellipse([cx - r + 2, cy - r + 2, cx + r - 2, cy + r - 2], 
+                                fill=(200, 200, 255, tint_alpha))
+            img = Image.alpha_composite(img.convert('RGBA'), overlay).convert('RGB')
+        
+        return np.array(img, dtype=np.float32)
+    
+    def _apply_dark_conditions(self, img, darkness_level):
+        """Simulate dark/low-light conditions with noise."""
+        # Reduce brightness
+        img = img * darkness_level
+        
+        # Add shot noise (Poisson-like) - more visible in dark
+        noise_scale = (1.0 - darkness_level) * 15
+        noise = self.rng.randn(*img.shape) * noise_scale
+        img = img + noise
+        
+        # Color temperature shift (warm/cool tint from artificial lighting)
+        if self.rng.random() > 0.5:
+            # Warm (yellowish - indoor lights)
+            img[:, :, 0] *= 1.1
+            img[:, :, 2] *= 0.85
+        else:
+            # Cool (bluish - screen light)
+            img[:, :, 0] *= 0.85
+            img[:, :, 2] *= 1.1
+        
+        return np.clip(img, 0, 255)
+    
+    def _apply_illumination_perturbation(self, img):
+        """Apply directional light gradient (from AGE framework)."""
+        size = img.shape[0]
+        
+        # Random gradient direction
+        angle = self.rng.random() * 2 * math.pi
+        
+        # Create gradient
+        y_coords, x_coords = np.mgrid[0:size, 0:size].astype(np.float32) / size
+        gradient = (x_coords * math.cos(angle) + y_coords * math.sin(angle))
+        gradient = (gradient - gradient.min()) / (gradient.max() - gradient.min() + 1e-8)
+        
+        # Random intensity and color
+        intensity = self.rng.uniform(0.1, 0.5)
+        color = self.rng.uniform(0.5, 1.5, size=3)
+        
+        gradient_rgb = np.stack([gradient * color[i] for i in range(3)], axis=-1)
+        
+        img = img + gradient_rgb * 255 * intensity
+        return np.clip(img, 0, 255)
+    
+    def _apply_sensor_noise(self, img):
+        """Simulate CMOS sensor noise (from AGE framework)."""
+        # Gaussian read noise
+        read_noise = self.rng.randn(*img.shape) * self.rng.uniform(2, 8)
+        # Shot noise (signal-dependent)
+        shot_noise = self.rng.randn(*img.shape) * np.sqrt(np.maximum(img, 0) + 1) * self.rng.uniform(0.1, 0.4)
+        # Fixed pattern noise
+        fpn = self.rng.randn(1, img.shape[1], img.shape[2]) * self.rng.uniform(1, 3)
+        
+        img = img + read_noise + shot_noise + fpn
+        return np.clip(img, 0, 255)
+    
+    def generate_sample(self, with_glasses_prob=0.25, dark_prob=0.3, 
+                        lazy_eye_prob=0.15, noise_prob=0.5):
+        """Generate a single training sample."""
+        # Random gaze target on screen
+        gaze_x = self.rng.uniform(0.05, 0.95)
+        gaze_y = self.rng.uniform(0.05, 0.95)
+        
+        # Random appearance
+        skin_tone = self.skin_tones[self.rng.randint(len(self.skin_tones))]
+        eye_color = self.eye_colors[self.rng.randint(len(self.eye_colors))]
+        eye_openness = self.rng.uniform(0.6, 1.0)
+        
+        # Lazy eye simulation: one eye deviates from the target
+        lazy_offset_x_L, lazy_offset_y_L = 0.0, 0.0
+        lazy_offset_x_R, lazy_offset_y_R = 0.0, 0.0
+        
+        if self.rng.random() < lazy_eye_prob:
+            # Strabismus: one eye deviates
+            affected_eye = self.rng.choice(['left', 'right'])
+            deviation_x = self.rng.uniform(-0.4, 0.4)
+            deviation_y = self.rng.uniform(-0.15, 0.15)
+            if affected_eye == 'left':
+                lazy_offset_x_L = deviation_x
+                lazy_offset_y_L = deviation_y
+            else:
+                lazy_offset_x_R = deviation_x
+                lazy_offset_y_R = deviation_y
+        
+        # Draw eyes
+        left_eye = self._draw_eye(gaze_x, gaze_y, skin_tone, eye_color, eye_openness,
+                                  lazy_offset_x_L, lazy_offset_y_L)
+        right_eye = self._draw_eye(gaze_x, gaze_y, skin_tone, eye_color, eye_openness,
+                                   lazy_offset_x_R, lazy_offset_y_R)
+        face = self._draw_face(skin_tone)
+        
+        # Apply glasses
+        if self.rng.random() < with_glasses_prob:
+            glasses_color = self.glasses_colors[self.rng.randint(len(self.glasses_colors))]
+            left_eye = self._add_glasses(left_eye, glasses_color)
+            right_eye = self._add_glasses(right_eye, glasses_color)
+        
+        # Apply dark conditions
+        if self.rng.random() < dark_prob:
+            darkness = self.rng.uniform(0.15, 0.5)
+            left_eye = self._apply_dark_conditions(left_eye, darkness)
+            right_eye = self._apply_dark_conditions(right_eye, darkness)
+            face = self._apply_dark_conditions(face, darkness)
+        
+        # Illumination perturbation
+        if self.rng.random() > 0.5:
+            left_eye = self._apply_illumination_perturbation(left_eye)
+            right_eye = self._apply_illumination_perturbation(right_eye)
+        
+        # Sensor noise
+        if self.rng.random() < noise_prob:
+            left_eye = self._apply_sensor_noise(left_eye)
+            right_eye = self._apply_sensor_noise(right_eye)
+        
+        # Normalize to [0, 1]
+        left_eye = left_eye / 255.0
+        right_eye = right_eye / 255.0
+        face = face / 255.0
+        
+        return {
+            'left_eye': left_eye.astype(np.float32),
+            'right_eye': right_eye.astype(np.float32),
+            'face': face.astype(np.float32),
+            'gaze_x': np.float32(gaze_x),
+            'gaze_y': np.float32(gaze_y),
+        }
+    
+    def generate_dataset(self, num_samples, with_glasses_prob=0.25, dark_prob=0.3,
+                         lazy_eye_prob=0.15):
+        """Generate a full dataset."""
+        left_eyes = []
+        right_eyes = []
+        faces = []
+        gaze_xs = []
+        gaze_ys = []
+        
+        for i in range(num_samples):
+            sample = self.generate_sample(
+                with_glasses_prob=with_glasses_prob,
+                dark_prob=dark_prob,
+                lazy_eye_prob=lazy_eye_prob
+            )
+            left_eyes.append(sample['left_eye'])
+            right_eyes.append(sample['right_eye'])
+            faces.append(sample['face'])
+            gaze_xs.append(sample['gaze_x'])
+            gaze_ys.append(sample['gaze_y'])
+            
+            if (i + 1) % 1000 == 0:
+                print(f"Generated {i+1}/{num_samples} samples")
+        
+        return {
+            'left_eye': np.array(left_eyes),
+            'right_eye': np.array(right_eyes),
+            'face': np.array(faces),
+            'gaze': np.column_stack([gaze_xs, gaze_ys])
+        }
+
+
+def create_tf_dataset(data_dict, batch_size=64, shuffle=True):
+    """Convert numpy arrays to tf.data.Dataset for training."""
+    dataset = tf.data.Dataset.from_tensor_slices((
+        {
+            'left_eye': data_dict['left_eye'],
+            'right_eye': data_dict['right_eye'],
+            'face': data_dict['face'],
+        },
+        data_dict['gaze']
+    ))
+    
+    if shuffle:
+        dataset = dataset.shuffle(buffer_size=min(len(data_dict['gaze']), 10000))
+    
+    dataset = dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)
+    return dataset
+
+
+def create_single_eye_dataset(data_dict, batch_size=64, shuffle=True):
+    """Create dataset for single-eye model (uses averaged eye features)."""
+    # Concatenate left and right eye side by side, or just use one
+    # For single-eye model, we combine both eye crops horizontally
+    # and also train on each eye separately for more data
+    
+    left_eyes = data_dict['left_eye']
+    right_eyes = data_dict['right_eye']
+    gaze = data_dict['gaze']
+    
+    # Use both eyes as separate training samples (doubles data)
+    all_eyes = np.concatenate([left_eyes, right_eyes], axis=0)
+    all_gaze = np.concatenate([gaze, gaze], axis=0)
+    
+    dataset = tf.data.Dataset.from_tensor_slices((all_eyes, all_gaze))
+    
+    if shuffle:
+        dataset = dataset.shuffle(buffer_size=min(len(all_gaze), 10000))
+    
+    dataset = dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)
+    return dataset
+
+
+if __name__ == '__main__':
+    print("Testing synthetic data generator...")
+    gen = SyntheticGazeDataGenerator(seed=42)
+    
+    # Generate a small batch
+    sample = gen.generate_sample()
+    print(f"Sample keys: {list(sample.keys())}")
+    print(f"Left eye shape: {sample['left_eye'].shape}")
+    print(f"Gaze: ({sample['gaze_x']:.3f}, {sample['gaze_y']:.3f})")
+    
+    # Generate small dataset
+    data = gen.generate_dataset(100)
+    print(f"\nDataset shapes:")
+    for k, v in data.items():
+        print(f"  {k}: {v.shape}")
+    
+    # Test tf.data pipeline
+    ds = create_tf_dataset(data, batch_size=16)
+    for inputs, labels in ds.take(1):
+        print(f"\nBatch shapes:")
+        for k, v in inputs.items():
+            print(f"  {k}: {v.shape}")
+        print(f"  labels: {labels.shape}")
+    
+    print("\nDone!")