fix bug when test load size larger tha img size

furyhawk · Sep 20, 2022 · f6eb2b2 · f6eb2b2
1 parent 0350f5f
commit f6eb2b2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/yolov6/data/datasets.py b/yolov6/data/datasets.py
@@ -116,7 +116,7 @@ def __getitem__(self, index):
             else:
                 img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
 
-            shapes = (h0, w0), ((h / h0, w / w0), pad)  # for COCO mAP rescaling
+            shapes = (h0, w0), ((h * ratio / h0, w * ratio / w0), pad)  # for COCO mAP rescaling
 
             labels = self.labels[index].copy()
             if labels.size:
@@ -626,4 +626,4 @@ def add_video(self, path):
         self.cap = cv2.VideoCapture(path)
         self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
     def __len__(self):
-        return self.nf  # number of files
+        return self.nf  # number of files