AutoShape PosixPath support (#4047)

* AutoShape PosixPath support

Usage example:

````python
from pathlib import Path

model = ...
file = Path('data/images/zidane.jpg')

results = model(file)
```

* Update common.py
This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-07-18 15:25:37 +02:00
1 parent dd62e2d05c
commit 9dd33fd20f
1 file changed
+5 -5
+5 -5
View File
@@ -1,7 +1,7 @@
# YOLOv5 common modules # YOLOv5 common modules
from copy import copy from copy import copy
from pathlib import Path from pathlib import Path, PosixPath
import math import math
import numpy as np import numpy as np
@@ -232,8 +232,8 @@ class AutoShape(nn.Module):
@torch.no_grad() @torch.no_grad()
def forward(self, imgs, size=640, augment=False, profile=False): def forward(self, imgs, size=640, augment=False, profile=False):
# Inference from various sources. For height=640, width=1280, RGB images example inputs are: # Inference from various sources. For height=640, width=1280, RGB images example inputs are:
# filename: imgs = 'data/images/zidane.jpg' # filename: imgs = 'data/images/zidane.jpg' # str or PosixPath
# URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg' # URI: = 'https://ultralytics.com/images/zidane.jpg'
# OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3) # OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)
# PIL: = Image.open('image.jpg') # HWC x(640,1280,3) # PIL: = Image.open('image.jpg') # HWC x(640,1280,3)
# numpy: = np.zeros((640,1280,3)) # HWC # numpy: = np.zeros((640,1280,3)) # HWC
@@ -251,8 +251,8 @@ class AutoShape(nn.Module):
shape0, shape1, files = [], [], [] # image and inference shapes, filenames shape0, shape1, files = [], [], [] # image and inference shapes, filenames
for i, im in enumerate(imgs): for i, im in enumerate(imgs):
f = f'image{i}' # filename f = f'image{i}' # filename
if isinstance(im, str): # filename or uri if isinstance(im, (str, PosixPath)): # filename or uri
im, f = Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im), im im, f = Image.open(requests.get(im, stream=True).raw if str(im).startswith('http') else im), im
im = np.asarray(exif_transpose(im)) im = np.asarray(exif_transpose(im))
elif isinstance(im, Image.Image): # PIL Image elif isinstance(im, Image.Image): # PIL Image
im, f = np.asarray(exif_transpose(im)), getattr(im, 'filename', f) or f im, f = np.asarray(exif_transpose(im)), getattr(im, 'filename', f) or f