WIP(mean): fix a bug when mask_background is False

deepem/loss/mean.py

@@ -128,7 +128,7 @@ def forward(
 
         # Recompute external matrix
         mext = self.compute_ext_matrix(ids, groups, self.recompute_ext, device)
-        vecs = self.generate_vecs(embd, trgt, ids)
+        vecs = self.generate_vecs(embd, trgt, mask, ids)
         means = [torch.mean(vec, dim=0) for vec in vecs]
         weights = [1.0] * len(vecs)
 
@@ -194,16 +194,21 @@ def generate_vecs(
         self,
         embd: torch.Tensor,
         trgt: torch.Tensor,
+        mask: torch.Tensor,
         ids: Sequence[int],
     ) -> list[torch.Tensor]:
         """
         Generate a list of vectorized embeddings for each ground truth object.
         """
+        mask_bool = mask.bool() if not self.mask_background else None
         result = []
         for obj_id in ids:
-            obj = torch.nonzero(trgt == int(obj_id))
-            z, y, x = obj[:, -3], obj[:, -2], obj[:, -1]
-            vec = embd[0, :, z, y, x].transpose(0, 1)  # Count x Dim
+            obj_mask = (trgt == int(obj_id)) & mask_bool if mask_bool is not None else (trgt == int(obj_id))
+            idx = torch.nonzero(obj_mask, as_tuple=True)
+            if idx[0].numel() == 0:
+                # If there are no indices for this ID, skip to the next one
+                continue
+            vec = embd[0, :, idx[-3], idx[-2], idx[-1]].transpose(0, 1)  # Count x Dim
             result.append(vec)
         return result