fix inference bug

hhaAndroid · Jan 4, 2024 · f6358c5 · f6358c5
1 parent 63150e6
commit f6358c5
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/mmdet/models/dense_heads/grounding_dino_head.py b/mmdet/models/dense_heads/grounding_dino_head.py
@@ -442,11 +442,13 @@ def _predict_by_feat_single(self,
             bbox_index = indexes // num_classes
             bbox_pred = bbox_pred[bbox_index]
         else:
+            # TODO: REC 
             cls_score = cls_score.sigmoid()
-            scores, _ = cls_score.max(-1)
-            scores, indexes = scores.topk(max_per_img)
-            bbox_pred = bbox_pred[indexes]
-            det_labels = scores.new_zeros(scores.shape, dtype=torch.long)
+            scores, indexes = cls_score.view(-1).topk(max_per_img)
+            num_classes = cls_score.shape[-1]
+            det_labels = indexes % num_classes
+            bbox_index = indexes // num_classes
+            bbox_pred = bbox_pred[bbox_index]
 
         det_bboxes = bbox_cxcywh_to_xyxy(bbox_pred)
         det_bboxes[:, 0::2] = det_bboxes[:, 0::2] * img_shape[1]

diff --git a/projects/mm_gdino_clip/batch_sampler.py b/projects/mm_gdino_clip/batch_sampler.py
@@ -36,7 +36,7 @@ def __iter__(self) -> Sequence[int]:
 
             od_to_rec_flag = False
             if dataset_mode == 'OD':
-                if np.random.random() > 1 - self.od_to_rec_prob:
+                if np.random.random() >= 1 - self.od_to_rec_prob:
                     dataset_mode = 'REC'
                     od_to_rec_flag = True
             else: