fix(data): fix a potential bug after closing the mosaic augmentation (#…

…603)
Megvii-BaseDetection · Aug 27, 2021 · 40bde39 · 40bde39
1 parent a1a2958
commit 40bde39
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/yolox/data/datasets/mosaicdetection.py b/yolox/data/datasets/mosaicdetection.py
@@ -154,9 +154,9 @@ def __getitem__(self, idx):
 
         else:
             self._dataset._input_dim = self.input_dim
-            img, label, img_info, id_ = self._dataset.pull_item(idx)
+            img, label, img_info, idx = self._dataset.pull_item(idx)
             img, label = self.preproc(img, label, self.input_dim)
-            return img, label, img_info, id_
+            return img, label, img_info, np.array([idx])
 
     def mixup(self, origin_img, origin_labels, input_dim):
         jit_factor = random.uniform(*self.mixup_scale)