Upload data/augmentations.py with huggingface_hub

data/augmentations.py

@@ -0,0 +1,309 @@
+"""
+Augmentation Pipeline for Face Detection.
+
+Implements SCRFD's "Sample Redistribution" strategy plus production-grade
+robustness augmentations for:
+- Tiny faces (large-scale crops generate small face positives)
+- Blur (Gaussian, motion blur)
+- Compression artifacts (JPEG quality degradation)
+- Low-light / poor illumination (brightness/gamma jitter)
+- Occlusion (random erasing simulating partial occlusion)
+
+Training augmentation pipeline (from SCRFD + TinaFace papers):
+1. Random crop with scale [0.3, 2.0] (Sample Redistribution)
+2. Resize to target size (640×640)
+3. Photometric distortion (brightness, contrast, hue, saturation)
+4. Horizontal flip (p=0.5)
+5. Random blur / compression / lighting degradation
+6. Normalize (ImageNet stats)
+"""
+
+import numpy as np
+import cv2
+from typing import Dict, Tuple, Optional
+
+
+class TrainAugmentation:
+    """
+    Full training augmentation with SCRFD Sample Redistribution.
+
+    The key insight: using crop scales up to 2.0× generates more
+    small-face positive anchors at stride 8 (72K → 118K per paper).
+    """
+
+    def __init__(self,
+                 target_size: int = 640,
+                 crop_scales: list = None,
+                 mean: tuple = (104.0, 117.0, 123.0),
+                 flip_prob: float = 0.5,
+                 enable_robustness: bool = True):
+        self.target_size = target_size
+        self.crop_scales = crop_scales or [0.3, 0.45, 0.6, 0.8, 1.0, 1.2, 1.4, 1.6, 1.8, 2.0]
+        self.mean = np.array(mean, dtype=np.float32)
+        self.flip_prob = flip_prob
+        self.enable_robustness = enable_robustness
+        self.robustness_aug = RobustnessAugmentation() if enable_robustness else None
+
+    def __call__(self, image: np.ndarray, boxes: np.ndarray,
+                 landmarks: np.ndarray) -> Dict:
+        h, w = image.shape[:2]
+
+        # 1. Random crop with Sample Redistribution
+        image, boxes, landmarks = self._random_crop(image, boxes, landmarks)
+
+        # 2. Resize to target
+        image, boxes, landmarks = self._resize(image, boxes, landmarks)
+
+        # 3. Photometric distortion
+        image = self._photometric_distort(image)
+
+        # 4. Horizontal flip
+        if np.random.random() < self.flip_prob:
+            image, boxes, landmarks = self._hflip(image, boxes, landmarks)
+
+        # 5. Robustness augmentations (blur, compression, lighting)
+        if self.enable_robustness and self.robustness_aug:
+            image = self.robustness_aug(image)
+
+        # 6. Mean subtraction (SCRFD-style normalization)
+        image = image.astype(np.float32) - self.mean
+
+        return {'image': image, 'boxes': boxes, 'landmarks': landmarks}
+
+    def _random_crop(self, image: np.ndarray, boxes: np.ndarray,
+                     landmarks: np.ndarray) -> Tuple:
+        """Random crop with sample redistribution scales."""
+        h, w = image.shape[:2]
+        scale = np.random.choice(self.crop_scales)
+        crop_size = int(min(h, w) * scale)
+        crop_size = max(crop_size, 32)
+
+        # If crop is larger than image, pad first
+        if crop_size > max(h, w):
+            pad_h = max(crop_size - h, 0)
+            pad_w = max(crop_size - w, 0)
+            image = cv2.copyMakeBorder(image, 0, pad_h, 0, pad_w,
+                                        cv2.BORDER_CONSTANT, value=(0, 0, 0))
+            h, w = image.shape[:2]
+
+        # Random crop location
+        max_x = w - crop_size
+        max_y = h - crop_size
+        x1 = np.random.randint(0, max(max_x, 1))
+        y1 = np.random.randint(0, max(max_y, 1))
+        x2 = x1 + crop_size
+        y2 = y1 + crop_size
+
+        # Crop image
+        cropped = image[y1:y2, x1:x2]
+
+        # Adjust boxes
+        new_boxes = boxes.copy()
+        new_boxes[:, 0] -= x1
+        new_boxes[:, 1] -= y1
+        new_boxes[:, 2] -= x1
+        new_boxes[:, 3] -= y1
+
+        # Clip to crop boundaries
+        new_boxes[:, 0] = np.clip(new_boxes[:, 0], 0, crop_size)
+        new_boxes[:, 1] = np.clip(new_boxes[:, 1], 0, crop_size)
+        new_boxes[:, 2] = np.clip(new_boxes[:, 2], 0, crop_size)
+        new_boxes[:, 3] = np.clip(new_boxes[:, 3], 0, crop_size)
+
+        # Filter valid boxes (at least 20% of original area visible)
+        orig_areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
+        new_widths = new_boxes[:, 2] - new_boxes[:, 0]
+        new_heights = new_boxes[:, 3] - new_boxes[:, 1]
+        new_areas = new_widths * new_heights
+        valid = (new_widths > 2) & (new_heights > 2) & (new_areas > 0.2 * orig_areas)
+
+        if valid.sum() == 0:
+            # Fallback: return original image
+            return image[:min(h, w), :min(h, w)], boxes, landmarks
+
+        new_boxes = new_boxes[valid]
+
+        # Adjust landmarks
+        new_lmk = landmarks[valid].copy()
+        for i in range(5):
+            new_lmk[:, i*2] -= x1
+            new_lmk[:, i*2+1] -= y1
+
+        return cropped, new_boxes, new_lmk
+
+    def _resize(self, image: np.ndarray, boxes: np.ndarray,
+                landmarks: np.ndarray) -> Tuple:
+        """Resize to target size."""
+        h, w = image.shape[:2]
+        scale_x = self.target_size / w
+        scale_y = self.target_size / h
+
+        image = cv2.resize(image, (self.target_size, self.target_size))
+
+        boxes[:, 0] *= scale_x
+        boxes[:, 1] *= scale_y
+        boxes[:, 2] *= scale_x
+        boxes[:, 3] *= scale_y
+
+        for i in range(5):
+            landmarks[:, i*2] *= scale_x
+            landmarks[:, i*2+1] *= scale_y
+
+        return image, boxes, landmarks
+
+    def _photometric_distort(self, image: np.ndarray) -> np.ndarray:
+        """Random photometric distortion (brightness, contrast, hue, saturation)."""
+        image = image.astype(np.float32)
+
+        # Brightness
+        if np.random.random() < 0.5:
+            delta = np.random.uniform(-32, 32)
+            image += delta
+
+        # Contrast
+        if np.random.random() < 0.5:
+            alpha = np.random.uniform(0.5, 1.5)
+            image *= alpha
+
+        # Color jitter in HSV
+        if np.random.random() < 0.5:
+            image_uint8 = np.clip(image, 0, 255).astype(np.uint8)
+            hsv = cv2.cvtColor(image_uint8, cv2.COLOR_RGB2HSV).astype(np.float32)
+
+            # Hue
+            hsv[:, :, 0] += np.random.uniform(-18, 18)
+            hsv[:, :, 0] = np.clip(hsv[:, :, 0], 0, 180)
+
+            # Saturation
+            hsv[:, :, 1] *= np.random.uniform(0.5, 1.5)
+            hsv[:, :, 1] = np.clip(hsv[:, :, 1], 0, 255)
+
+            image = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2RGB).astype(np.float32)
+
+        return np.clip(image, 0, 255)
+
+    def _hflip(self, image: np.ndarray, boxes: np.ndarray,
+               landmarks: np.ndarray) -> Tuple:
+        """Horizontal flip with landmark reordering."""
+        w = image.shape[1]
+        image = image[:, ::-1].copy()
+
+        new_boxes = boxes.copy()
+        new_boxes[:, 0] = w - boxes[:, 2]
+        new_boxes[:, 2] = w - boxes[:, 0]
+
+        new_lmk = landmarks.copy()
+        for i in range(5):
+            new_lmk[:, i*2] = w - landmarks[:, i*2]
+
+        # Reorder landmarks for face symmetry:
+        # Standard 5-point: left_eye, right_eye, nose, left_mouth, right_mouth
+        # After flip: swap left↔right
+        if new_lmk.shape[0] > 0 and np.any(new_lmk > 0):
+            # Swap left_eye ↔ right_eye
+            new_lmk[:, [0, 1, 2, 3]] = new_lmk[:, [2, 3, 0, 1]]
+            # Swap left_mouth ↔ right_mouth
+            new_lmk[:, [6, 7, 8, 9]] = new_lmk[:, [8, 9, 6, 7]]
+
+        return image, new_boxes, new_lmk
+
+
+class ValAugmentation:
+    """Validation: resize + normalize only."""
+
+    def __init__(self, target_size: int = 640,
+                 mean: tuple = (104.0, 117.0, 123.0)):
+        self.target_size = target_size
+        self.mean = np.array(mean, dtype=np.float32)
+
+    def __call__(self, image: np.ndarray, boxes: np.ndarray,
+                 landmarks: np.ndarray) -> Dict:
+        h, w = image.shape[:2]
+
+        # Resize keeping aspect ratio
+        scale = self.target_size / max(h, w)
+        new_h, new_w = int(h * scale), int(w * scale)
+        image = cv2.resize(image, (new_w, new_h))
+
+        # Pad to target size
+        pad_h = self.target_size - new_h
+        pad_w = self.target_size - new_w
+        image = cv2.copyMakeBorder(image, 0, pad_h, 0, pad_w,
+                                    cv2.BORDER_CONSTANT, value=(0, 0, 0))
+
+        # Scale boxes
+        boxes[:, 0] *= scale
+        boxes[:, 1] *= scale
+        boxes[:, 2] *= scale
+        boxes[:, 3] *= scale
+
+        for i in range(5):
+            landmarks[:, i*2] *= scale
+            landmarks[:, i*2+1] *= scale
+
+        image = image.astype(np.float32) - self.mean
+
+        return {'image': image, 'boxes': boxes, 'landmarks': landmarks}
+
+
+class RobustnessAugmentation:
+    """
+    Production-grade robustness augmentations targeting known failure modes.
+
+    Applied with probability during training to make the detector robust to:
+    1. Gaussian blur (σ = 0.5–3.0) — camera defocus, motion blur
+    2. JPEG compression (Q = 20–80) — streaming/compression artifacts
+    3. Low-light gamma (γ = 1.5–3.0) — dark environments
+    4. Random occlusion (Cutout) — partial face occlusion
+    5. Gaussian noise — sensor noise, low-light grain
+    """
+
+    def __init__(self,
+                 blur_prob: float = 0.2,
+                 jpeg_prob: float = 0.2,
+                 lowlight_prob: float = 0.15,
+                 occlusion_prob: float = 0.1,
+                 noise_prob: float = 0.15):
+        self.blur_prob = blur_prob
+        self.jpeg_prob = jpeg_prob
+        self.lowlight_prob = lowlight_prob
+        self.occlusion_prob = occlusion_prob
+        self.noise_prob = noise_prob
+
+    def __call__(self, image: np.ndarray) -> np.ndarray:
+        # Gaussian blur
+        if np.random.random() < self.blur_prob:
+            sigma = np.random.uniform(0.5, 3.0)
+            ksize = int(sigma * 6) | 1  # Ensure odd
+            image = cv2.GaussianBlur(image, (ksize, ksize), sigma)
+
+        # JPEG compression artifacts
+        if np.random.random() < self.jpeg_prob:
+            quality = np.random.randint(20, 80)
+            encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
+            _, buf = cv2.imencode('.jpg', image.astype(np.uint8), encode_param)
+            image = cv2.imdecode(buf, cv2.IMREAD_COLOR).astype(np.float32)
+
+        # Low-light simulation (gamma darkening)
+        if np.random.random() < self.lowlight_prob:
+            gamma = np.random.uniform(1.5, 3.0)
+            image = np.clip(image, 0, 255)
+            image = ((image / 255.0) ** gamma * 255.0)
+
+        # Random occlusion (Cutout)
+        if np.random.random() < self.occlusion_prob:
+            h, w = image.shape[:2]
+            # Random rectangle
+            rh = np.random.randint(h // 10, h // 4)
+            rw = np.random.randint(w // 10, w // 4)
+            ry = np.random.randint(0, h - rh)
+            rx = np.random.randint(0, w - rw)
+            image[ry:ry+rh, rx:rx+rw] = np.random.randint(0, 255, 3)
+
+        # Gaussian noise
+        if np.random.random() < self.noise_prob:
+            sigma = np.random.uniform(5, 25)
+            noise = np.random.randn(*image.shape) * sigma
+            image = np.clip(image + noise, 0, 255)
+
+        return image.astype(np.float32)