Bugfix: LoFTR was ignoring mask input (#2923)

kornia · Jun 10, 2024 · e17ce88 · e17ce88
1 parent 1e0e648
commit e17ce88
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/kornia/feature/loftr/loftr.py b/kornia/feature/loftr/loftr.py
@@ -154,9 +154,9 @@ def forward(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
         feat_c1 = feat_c1.reshape(n1, -1, c1)
 
         mask_c0 = mask_c1 = None  # mask is useful in training
-        if "mask0" in _data:
+        if "mask0" in data:
             mask_c0 = resize(data["mask0"], _data["hw0_c"], interpolation="nearest").flatten(-2)
-        if "mask1" in _data:
+        if "mask1" in data:
             mask_c1 = resize(data["mask1"], _data["hw1_c"], interpolation="nearest").flatten(-2)
         feat_c0, feat_c1 = self.loftr_coarse(feat_c0, feat_c1, mask_c0, mask_c1)