Update landmarkdiff/synthetic/augmentation.py to v0.3.2

landmarkdiff/synthetic/augmentation.py

@@ -1,8 +1,9 @@
-"""Clinical degradation augmentations.
+"""Clinical degradation augmentation pipeline.
 
-Degrades clean FFHQ/CelebA-HQ to match real clinical photo distribution.
-Applied from day 1 - domain gap prevention, not afterthought.
-3-5 random augmentations per sample.
+Degrades clean FFHQ/CelebA-HQ images to match real clinical photo distribution.
+Applied from day 1 of training — domain gap prevention, not afterthought.
+
+Each sample gets 3-5 random augmentations from the pool.
 """
 
 from __future__ import annotations
@@ -26,6 +27,8 @@ class AugmentationConfig:
 def point_source_lighting(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
     """Simulate point-source clinical lighting from a random direction."""
     h, w = image.shape[:2]
+    if h < 4 or w < 4:
+        return image
 
     # Random light source position
     lx = rng.uniform(0, w)
@@ -35,7 +38,7 @@ def point_source_lighting(image: np.ndarray, rng: np.random.Generator) -> np.nda
     # Distance-based falloff
     y_grid, x_grid = np.mgrid[0:h, 0:w].astype(np.float32)
     dist = np.sqrt((x_grid - lx) ** 2 + (y_grid - ly) ** 2)
-    max_dist = np.sqrt(w**2 + h**2)
+    max_dist = np.sqrt(w ** 2 + h ** 2)
     light_map = 1.0 - (dist / max_dist) * intensity
 
     light_map = np.clip(light_map, 0.3, 1.0)
@@ -76,7 +79,8 @@ def jpeg_compression(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
     quality = int(rng.uniform(40, 85))
     encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]
     _, encoded = cv2.imencode(".jpg", image, encode_param)
-    return cv2.imdecode(encoded, cv2.IMREAD_COLOR)
+    decoded = cv2.imdecode(encoded, cv2.IMREAD_COLOR)
+    return decoded if decoded is not None else image
 
 
 def gaussian_sensor_noise(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
@@ -89,6 +93,8 @@ def gaussian_sensor_noise(image: np.ndarray, rng: np.random.Generator) -> np.nda
 def barrel_distortion(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
     """Apply barrel/pincushion distortion simulating phone camera lens."""
     h, w = image.shape[:2]
+    if h < 4 or w < 4:
+        return image
     k1 = rng.uniform(-0.2, 0.2)
 
     fx = fy = max(w, h)
@@ -105,6 +111,9 @@ def barrel_distortion(image: np.ndarray, rng: np.random.Generator) -> np.ndarray
 
 def motion_blur(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
     """Slight motion blur (common in handheld clinical photos)."""
+    h, w = image.shape[:2]
+    if h < 4 or w < 4:
+        return image
     size = int(rng.uniform(3, 7))
     angle = rng.uniform(0, 180)
 
@@ -117,7 +126,6 @@ def motion_blur(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
     if ksum > 0:
         kernel = kernel / ksum
     else:
-        # rotation can zero out the kernel - fall back to identity
         kernel = np.zeros_like(kernel)
         kernel[size // 2, size // 2] = 1.0
 
@@ -127,12 +135,14 @@ def motion_blur(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
 def vignette(image: np.ndarray, rng: np.random.Generator) -> np.ndarray:
     """Add lens vignetting (darkened corners)."""
     h, w = image.shape[:2]
+    if h < 4 or w < 4:
+        return image
     strength = rng.uniform(0.3, 0.7)
 
     y, x = np.mgrid[0:h, 0:w].astype(np.float32)
     cx, cy = w / 2, h / 2
     dist = np.sqrt((x - cx) ** 2 + (y - cy) ** 2)
-    max_dist = np.sqrt(cx**2 + cy**2)
+    max_dist = np.sqrt(cx ** 2 + cy ** 2)
 
     mask = 1 - strength * (dist / max_dist) ** 2
     mask = np.clip(mask, 0.3, 1.0)
@@ -160,7 +170,20 @@ def apply_clinical_augmentation(
     max_augmentations: int = 5,
     rng: np.random.Generator | None = None,
 ) -> np.ndarray:
-    """Apply random clinical degradation augmentations to an image."""
+    """Apply random clinical degradation augmentations to an image.
+
+    Each sample gets min_augmentations to max_augmentations from the pool,
+    selected by their individual probabilities.
+
+    Args:
+        image: BGR input image (clean FFHQ/CelebA-HQ).
+        min_augmentations: Minimum number of augmentations to apply.
+        max_augmentations: Maximum number of augmentations to apply.
+        rng: Random number generator.
+
+    Returns:
+        Degraded image matching clinical photo distribution.
+    """
     rng = rng or np.random.default_rng()
 
     # Select augmentations by probability