Merge pull request #376 from WongGawa/cls_reshape

fix difference bewtween high version albumentations
mindspore-lab · Oct 18, 2024 · dfe6041 · dfe6041
2 parents 5e86022 + af27b75
commit dfe6041
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mindyolo/data/albumentations.py b/mindyolo/data/albumentations.py
@@ -55,8 +55,8 @@ def __call__(self, sample, p=1.0, **kwargs):
             new = self.transform(image=im, bboxes=bboxes, class_labels=cls)  # transformed
 
             sample['img'] = new['image']
-            sample['bboxes'] = np.array(new['bboxes']).reshape(-1, 1)
-            sample['cls'] = np.array(new['class_labels'])
+            sample['bboxes'] = np.array(new['bboxes'])
+            sample['cls'] = np.array(new['class_labels']).reshape(-1, 1)
             sample['bbox_format'] = "xywhn"
 
         return sample