Update landmarkdiff/losses.py to v0.3.2

landmarkdiff/losses.py

@@ -1,7 +1,6 @@
 """4-term loss function module for ControlNet fine-tuning.
 
-L_total = L_diffusion + w_landmark * L_landmark
-        + w_identity * L_identity + w_perceptual * L_perceptual
+L_total = L_diffusion + w_landmark * L_landmark + w_identity * L_identity + w_perceptual * L_perceptual
 
 Phase A (synthetic TPS data): L_diffusion ONLY. No perceptual loss against
 rubbery TPS warps — it would penalize realism.
@@ -23,8 +22,8 @@ class LossWeights:
 
     diffusion: float = 1.0
     landmark: float = 0.1
-    identity: float = 0.05
-    perceptual: float = 0.1
+    identity: float = 0.1
+    perceptual: float = 0.05
 
 
 class DiffusionLoss:
@@ -93,16 +92,11 @@ class IdentityLoss:
             return
         try:
             from insightface.app import FaceAnalysis
-
             self._app = FaceAnalysis(
                 name="buffalo_l",
                 providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
             )
-            ctx_id = (
-                device.index
-                if device.type == "cuda" and device.index is not None
-                else (0 if device.type == "cuda" else -1)
-            )
+            ctx_id = device.index if device.type == "cuda" and device.index is not None else (0 if device.type == "cuda" else -1)
             self._app.prepare(ctx_id=ctx_id, det_size=(320, 320))
             self._has_arcface = True
         except Exception:
@@ -120,7 +114,6 @@ class IdentityLoss:
         """
         if self._has_arcface:
             import numpy as np
-
             embeddings = []
             valid_mask = []
             for i in range(image_tensor.shape[0]):
@@ -159,9 +152,7 @@ class IdentityLoss:
 
         # Resize to 112x112 for ArcFace
         pred_112 = F.interpolate(pred_crop, size=(112, 112), mode="bilinear", align_corners=False)
-        target_112 = F.interpolate(
-            target_crop, size=(112, 112), mode="bilinear", align_corners=False
-        )
+        target_112 = F.interpolate(target_crop, size=(112, 112), mode="bilinear", align_corners=False)
 
         # Normalize to [-1, 1]
         pred_norm = pred_112 * 2 - 1
@@ -172,7 +163,7 @@ class IdentityLoss:
         target_emb, target_valid = self._extract_embedding(target_norm)
 
         # Only compute loss for samples where both faces were detected
-        valid = [p and t for p, t in zip(pred_valid, target_valid, strict=False)]
+        valid = [p and t for p, t in zip(pred_valid, target_valid)]
         if not any(valid):
             return torch.tensor(0.0, device=pred_image.device)
 
@@ -225,7 +216,6 @@ class PerceptualLoss:
         if self._lpips is None:
             try:
                 import lpips
-
                 self._lpips = lpips.LPIPS(net="alex").to(device)
                 self._lpips.eval()
                 for p in self._lpips.parameters():
@@ -235,33 +225,31 @@ class PerceptualLoss:
 
     def __call__(
         self,
-        pred: torch.Tensor,  # (B, 3, H, W) in [0, 1]
+        pred: torch.Tensor,    # (B, 3, H, W) in [0, 1]
         target: torch.Tensor,
-        mask: torch.Tensor,  # (B, 1, H, W) surgical mask [0, 1]
+        mask: torch.Tensor,    # (B, 1, H, W) surgical mask [0, 1]
     ) -> torch.Tensor:
         self._ensure_loaded(pred.device)
 
-        # Invert mask: we want loss OUTSIDE surgical region
-        outside_mask = 1 - mask
-
-        # Erode outside_mask to exclude boundary pixels — avoids artificial
-        # edge features where masked (0) meets unmasked (non-zero) values
-        erode_kernel = 5
-        if outside_mask.shape[-1] >= erode_kernel and outside_mask.shape[-2] >= erode_kernel:
-            outside_mask = -F.max_pool2d(
-                -outside_mask,
-                kernel_size=erode_kernel,
-                stride=1,
-                padding=erode_kernel // 2,
-            )
-
-        # Normalize to [-1, 1] for LPIPS FIRST, then mask
+        # Normalize to [-1, 1] for LPIPS
         pred_norm = pred * 2 - 1
         target_norm = target * 2 - 1
 
-        # Apply mask after normalization (masked regions become 0, not -1)
-        pred_norm = pred_norm * outside_mask
-        target_norm = target_norm * outside_mask
+        # When mask is all-ones (no mask file available), compute on full image.
+        # Otherwise invert mask to get loss OUTSIDE the surgical region only.
+        has_mask = mask.sum() < mask.numel() * 0.99
+        if has_mask:
+            outside_mask = 1 - mask
+            erode_kernel = 5
+            if outside_mask.shape[-1] >= erode_kernel and outside_mask.shape[-2] >= erode_kernel:
+                outside_mask = -F.max_pool2d(
+                    -outside_mask,
+                    kernel_size=erode_kernel,
+                    stride=1,
+                    padding=erode_kernel // 2,
+                )
+            pred_norm = pred_norm * outside_mask
+            target_norm = target_norm * outside_mask
 
         if self._lpips == "unavailable":
             # Fallback: simple L1 loss
@@ -299,7 +287,6 @@ class CombinedLoss:
         # or ONNX-based fallback
         if use_differentiable_arcface:
             from landmarkdiff.arcface_torch import ArcFaceLoss
-
             self.identity_loss = ArcFaceLoss(weights_path=arcface_weights_path)
         else:
             self.identity_loss = IdentityLoss()