Update iou_neg_piecewise_sampler.py (open-mmlab#2017)

fix lint
ZwwWayne · Dec 3, 2022 · 3b0ae48 · 3b0ae48
1 parent aa544d4
commit 3b0ae48
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/mmdet3d/models/task_modules/samplers/iou_neg_piecewise_sampler.py b/mmdet3d/models/task_modules/samplers/iou_neg_piecewise_sampler.py
@@ -1,4 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import math
+
 import torch
 
 from mmdet3d.registry import TASK_UTILS
@@ -74,9 +76,11 @@ def _sample_neg(self, assign_result, num_expected, **kwargs):
                     # if the numbers of negative samplers in previous
                     # pieces are less than the expected number, extend
                     # the same number in the current piece.
-                    piece_expected_num = int(
-                        num_expected *
-                        self.neg_piece_fractions[piece_inds]) + extend_num
+                    piece_expected_num = min(
+                        num_expected,
+                        math.ceil(num_expected *
+                                  self.neg_piece_fractions[piece_inds]) +
+                        extend_num)
                     min_iou_thr = self.neg_iou_thr[piece_inds + 1]
                 max_iou_thr = self.neg_iou_thr[piece_inds]
                 piece_neg_inds = torch.nonzero(