Faster batch `same_shapes` (#8851)

9 months ago · 617ab21523
parent 2ea6b2b889
commit 617ab21523
1 changed files with 1 additions and 1 deletions
--- a/ultralytics/engine/predictor.py
+++ b/ultralytics/engine/predictor.py
@ -151,7 +151,7 @@ class BasePredictor:
        Returns:
            (list): A list of transformed images.
        """
-        same_shapes = all(x.shape == im[0].shape for x in im)
+        same_shapes = len({x.shape for x in im}) == 1
        letterbox = LetterBox(self.imgsz, auto=same_shapes and self.model.pt, stride=self.model.stride)
        return [letterbox(image=x) for x in im]