Update landmarkdiff/augmentation.py to v0.3.2

landmarkdiff/augmentation.py

@@ -21,7 +21,6 @@ import numpy as np
 @dataclass
 class AugmentationConfig:
     """Augmentation parameters."""
-
     # Geometric
     random_flip: bool = True
     random_rotation_deg: float = 5.0
@@ -92,9 +91,8 @@ def augment_training_sample(
         out_input = np.ascontiguousarray(out_input[:, ::-1])
         out_target = np.ascontiguousarray(out_target[:, ::-1])
         out_cond = np.ascontiguousarray(out_cond[:, ::-1])
-        out_mask = np.ascontiguousarray(
-            out_mask[:, ::-1] if out_mask.ndim == 2 else out_mask[:, ::-1, :]
-        )
+        out_mask = np.ascontiguousarray(out_mask[:, ::-1] if out_mask.ndim == 2
+                                        else out_mask[:, ::-1, :])
         if out_lm_src is not None:
             out_lm_src[:, 0] = 1.0 - out_lm_src[:, 0]
         if out_lm_dst is not None:
@@ -112,13 +110,15 @@ def augment_training_sample(
         M[0, 2] += tx
         M[1, 2] += ty
 
-        out_input = cv2.warpAffine(out_input, M, (w, h), borderMode=cv2.BORDER_REFLECT_101)
-        out_target = cv2.warpAffine(out_target, M, (w, h), borderMode=cv2.BORDER_REFLECT_101)
-        out_cond = cv2.warpAffine(
-            out_cond, M, (w, h), borderMode=cv2.BORDER_CONSTANT, borderValue=0
-        )
+        out_input = cv2.warpAffine(out_input, M, (w, h),
+                                    borderMode=cv2.BORDER_REFLECT_101)
+        out_target = cv2.warpAffine(out_target, M, (w, h),
+                                     borderMode=cv2.BORDER_REFLECT_101)
+        out_cond = cv2.warpAffine(out_cond, M, (w, h),
+                                   borderMode=cv2.BORDER_CONSTANT, borderValue=0)
         mask_2d = out_mask if out_mask.ndim == 2 else out_mask[:, :, 0]
-        mask_2d = cv2.warpAffine(mask_2d, M, (w, h), borderMode=cv2.BORDER_CONSTANT, borderValue=0)
+        mask_2d = cv2.warpAffine(mask_2d, M, (w, h),
+                                  borderMode=cv2.BORDER_CONSTANT, borderValue=0)
         out_mask = mask_2d if out_mask.ndim == 2 else mask_2d[:, :, np.newaxis]
 
         # Transform landmarks
@@ -183,7 +183,9 @@ def augment_training_sample(
     return result
 
 
-def _transform_landmarks(landmarks: np.ndarray, M: np.ndarray, w: int, h: int) -> np.ndarray:
+def _transform_landmarks(
+    landmarks: np.ndarray, M: np.ndarray, w: int, h: int
+) -> np.ndarray:
     """Transform normalized landmarks with an affine matrix."""
     # Convert to pixel coords
     px = landmarks.copy()
@@ -260,12 +262,8 @@ class FitzpatrickBalancer:
             target_distribution: Target fraction per type. Defaults to uniform.
         """
         self.target = target_distribution or {
-            "I": 1 / 6,
-            "II": 1 / 6,
-            "III": 1 / 6,
-            "IV": 1 / 6,
-            "V": 1 / 6,
-            "VI": 1 / 6,
+            "I": 1/6, "II": 1/6, "III": 1/6,
+            "IV": 1/6, "V": 1/6, "VI": 1/6,
         }
         self._counts: dict[str, int] = {}
 
@@ -284,7 +282,7 @@ class FitzpatrickBalancer:
         for ft in fitz_types:
             count = self._counts.get(ft, 1)
             freq = count / total
-            target_freq = self.target.get(ft, 1 / 6)
+            target_freq = self.target.get(ft, 1/6)
             # Weight = target / actual (capped for stability)
             w = min(target_freq / max(freq, 1e-6), 5.0)
             weights.append(w)