Following ultralytics/yolov5#5694

Author: zhiqwang

yolort/v5/helper.py

@@ -4,6 +4,7 @@
 
 import torch
 
+from .models import AutoShape
 from .models.yolo import Model
 from .utils import attempt_download, intersect_dicts, set_logging
 
@@ -68,6 +69,6 @@ def load_yolov5_model(checkpoint_path: str, autoshape: bool = False, verbose: bo
     model.load_state_dict(ckpt_state_dict, strict=False)
 
     if autoshape:
-        model = model.autoshape()
+        model = AutoShape(model)
 
     return model

yolort/v5/models/common.py

@@ -434,17 +434,17 @@ def autoshape(self):
 
     @torch.no_grad()
     def forward(self, imgs, size=640, augment=False, profile=False):
+        """
+        Inference from various sources. For height=640, width=1280, RGB images example inputs are:
+            - file: imgs = 'data/images/zidane.jpg'  # str or PosixPath
+            - URI: = 'https://ultralytics.com/images/zidane.jpg'
+            - OpenCV: = cv2.imread('image.jpg')[:,:,::-1]  # HWC BGR to RGB x(640,1280,3)
+            - PIL: = Image.open('image.jpg') or ImageGrab.grab()  # HWC x(640,1280,3)
+            - numpy: = np.zeros((640,1280,3))  # HWC
+            - torch: = torch.zeros(16,3,320,640)  # BCHW (scaled to size=640, 0-1 values)
+            - multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...]  # list of images
+        """
         from yolort.v5.utils.augmentations import letterbox
-
-        # Inference from various sources. For height=640, width=1280, RGB images example inputs are:
-        #   file:       imgs = 'data/images/zidane.jpg'  # str or PosixPath
-        #   URI:             = 'https://ultralytics.com/images/zidane.jpg'
-        #   OpenCV:          = cv2.imread('image.jpg')[:,:,::-1]  # HWC BGR to RGB x(640,1280,3)
-        #   PIL:             = Image.open('image.jpg') or ImageGrab.grab()  # HWC x(640,1280,3)
-        #   numpy:           = np.zeros((640,1280,3))  # HWC
-        #   torch:           = torch.zeros(16,3,320,640)  # BCHW (scaled to size=640, 0-1 values)
-        #   multiple:        = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...]  # list of images
-
         from yolort.v5.utils.datasets import exif_transpose
 
         t = [time_sync()]
@@ -454,17 +454,14 @@ def forward(self, imgs, size=640, augment=False, profile=False):
                 return self.model(imgs.to(p.device).type_as(p), augment, profile)  # inference
 
         # Pre-process
-        n, imgs = (
-            (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs])
-        )  # number of images, list of images
+        # number of images, list of images
+        n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs])
         shape0, shape1, files = [], [], []  # image and inference shapes, filenames
         for i, im in enumerate(imgs):
             f = f"image{i}"  # filename
             if isinstance(im, (str, Path)):  # filename or uri
-                im, f = (
-                    Image.open(requests.get(im, stream=True).raw if str(im).startswith("http") else im),
-                    im,
-                )
+                f = im
+                im = Image.open(requests.get(im, stream=True).raw if str(im).startswith("http") else im)
                 im = np.asarray(exif_transpose(im))
             elif isinstance(im, Image.Image):  # PIL Image
                 im, f = np.asarray(exif_transpose(im)), getattr(im, "filename", f) or f