From e9146c8144dea5f3cbdd98b225f3e147a171c216 Mon Sep 17 00:00:00 2001 From: wangyuang Date: Thu, 6 May 2021 23:03:00 +0800 Subject: [PATCH] FIX mask_border bug --- src/loftr/utils/coarse_matching.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/loftr/utils/coarse_matching.py b/src/loftr/utils/coarse_matching.py index 0570565..ffa8bfa 100644 --- a/src/loftr/utils/coarse_matching.py +++ b/src/loftr/utils/coarse_matching.py @@ -11,17 +11,23 @@ def mask_border(m, b: int, v): b (int) v (m.dtype) """ + if b <= 0: + return + m[:, :b] = v m[:, :, :b] = v m[:, :, :, :b] = v m[:, :, :, :, :b] = v - m[:, -b:0] = v - m[:, :, -b:0] = v - m[:, :, :, -b:0] = v - m[:, :, :, :, -b:0] = v + m[:, -b:] = v + m[:, :, -b:] = v + m[:, :, :, -b:] = v + m[:, :, :, :, -b:] = v def mask_border_with_padding(m, bd, v, p_m0, p_m1): + if bd <= 0: + return + m[:, :bd] = v m[:, :, :bd] = v m[:, :, :, :bd] = v