glenn-jocher
commited on
PyTorch Hub cv2 .save() .show() bug fix (#2831)
Browse files* PyTorch Hub cv2 .save() .show() bug fix
cv2.rectangle() was failing on non-contiguous np array inputs. This checks for contiguous arrays and applies is necessary:
```python
imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update
```
* Update plots.py
```python
assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to plot_on_box() input image.'
```
* Update hubconf.py
Expand CI tests to OpenCV image.
- hubconf.py +6 -4
- models/common.py +2 -2
- utils/plots.py +14 -12
hubconf.py
CHANGED
@@ -124,13 +124,15 @@ if __name__ == '__main__':
|
|
124 |
# model = custom(path_or_model='path/to/model.pt') # custom example
|
125 |
|
126 |
# Verify inference
|
|
|
127 |
import numpy as np
|
128 |
from PIL import Image
|
129 |
|
130 |
-
imgs = [
|
131 |
-
'
|
132 |
-
'
|
133 |
-
|
|
|
134 |
|
135 |
results = model(imgs) # batched inference
|
136 |
results.print()
|
|
|
124 |
# model = custom(path_or_model='path/to/model.pt') # custom example
|
125 |
|
126 |
# Verify inference
|
127 |
+
import cv2
|
128 |
import numpy as np
|
129 |
from PIL import Image
|
130 |
|
131 |
+
imgs = ['data/images/zidane.jpg', # filename
|
132 |
+
'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg', # URI
|
133 |
+
cv2.imread('data/images/bus.jpg')[:, :, ::-1], # OpenCV
|
134 |
+
Image.open('data/images/bus.jpg'), # PIL
|
135 |
+
np.zeros((320, 640, 3))] # numpy
|
136 |
|
137 |
results = model(imgs) # batched inference
|
138 |
results.print()
|
models/common.py
CHANGED
@@ -240,7 +240,7 @@ class autoShape(nn.Module):
|
|
240 |
@torch.no_grad()
|
241 |
def forward(self, imgs, size=640, augment=False, profile=False):
|
242 |
# Inference from various sources. For height=640, width=1280, RGB images example inputs are:
|
243 |
-
# filename: imgs = 'data/
|
244 |
# URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg'
|
245 |
# OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)
|
246 |
# PIL: = Image.open('image.jpg') # HWC x(640,1280,3)
|
@@ -271,7 +271,7 @@ class autoShape(nn.Module):
|
|
271 |
shape0.append(s) # image shape
|
272 |
g = (size / max(s)) # gain
|
273 |
shape1.append([y * g for y in s])
|
274 |
-
imgs[i] = im # update
|
275 |
shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)] # inference shape
|
276 |
x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad
|
277 |
x = np.stack(x, 0) if n > 1 else x[0][None] # stack
|
|
|
240 |
@torch.no_grad()
|
241 |
def forward(self, imgs, size=640, augment=False, profile=False):
|
242 |
# Inference from various sources. For height=640, width=1280, RGB images example inputs are:
|
243 |
+
# filename: imgs = 'data/images/zidane.jpg'
|
244 |
# URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg'
|
245 |
# OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)
|
246 |
# PIL: = Image.open('image.jpg') # HWC x(640,1280,3)
|
|
|
271 |
shape0.append(s) # image shape
|
272 |
g = (size / max(s)) # gain
|
273 |
shape1.append([y * g for y in s])
|
274 |
+
imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update
|
275 |
shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)] # inference shape
|
276 |
x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad
|
277 |
x = np.stack(x, 0) if n > 1 else x[0][None] # stack
|
utils/plots.py
CHANGED
@@ -54,32 +54,34 @@ def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5):
|
|
54 |
return filtfilt(b, a, data) # forward-backward filter
|
55 |
|
56 |
|
57 |
-
def plot_one_box(x,
|
58 |
-
# Plots one bounding box on image
|
59 |
-
|
|
|
60 |
color = color or [random.randint(0, 255) for _ in range(3)]
|
61 |
c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
|
62 |
-
cv2.rectangle(
|
63 |
if label:
|
64 |
tf = max(tl - 1, 1) # font thickness
|
65 |
t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
|
66 |
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
|
67 |
-
cv2.rectangle(
|
68 |
-
cv2.putText(
|
69 |
|
70 |
|
71 |
-
def plot_one_box_PIL(box,
|
72 |
-
|
73 |
-
|
74 |
-
|
|
|
75 |
draw.rectangle(box, width=line_thickness, outline=tuple(color)) # plot
|
76 |
if label:
|
77 |
-
fontsize = max(round(max(
|
78 |
font = ImageFont.truetype("Arial.ttf", fontsize)
|
79 |
txt_width, txt_height = font.getsize(label)
|
80 |
draw.rectangle([box[0], box[1] - txt_height + 4, box[0] + txt_width, box[1]], fill=tuple(color))
|
81 |
draw.text((box[0], box[1] - txt_height + 1), label, fill=(255, 255, 255), font=font)
|
82 |
-
return np.asarray(
|
83 |
|
84 |
|
85 |
def plot_wh_methods(): # from utils.plots import *; plot_wh_methods()
|
|
|
54 |
return filtfilt(b, a, data) # forward-backward filter
|
55 |
|
56 |
|
57 |
+
def plot_one_box(x, im, color=None, label=None, line_thickness=3):
|
58 |
+
# Plots one bounding box on image 'im' using OpenCV
|
59 |
+
assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to plot_on_box() input image.'
|
60 |
+
tl = line_thickness or round(0.002 * (im.shape[0] + im.shape[1]) / 2) + 1 # line/font thickness
|
61 |
color = color or [random.randint(0, 255) for _ in range(3)]
|
62 |
c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
|
63 |
+
cv2.rectangle(im, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
|
64 |
if label:
|
65 |
tf = max(tl - 1, 1) # font thickness
|
66 |
t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
|
67 |
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
|
68 |
+
cv2.rectangle(im, c1, c2, color, -1, cv2.LINE_AA) # filled
|
69 |
+
cv2.putText(im, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
|
70 |
|
71 |
|
72 |
+
def plot_one_box_PIL(box, im, color=None, label=None, line_thickness=None):
|
73 |
+
# Plots one bounding box on image 'im' using PIL
|
74 |
+
im = Image.fromarray(im)
|
75 |
+
draw = ImageDraw.Draw(im)
|
76 |
+
line_thickness = line_thickness or max(int(min(im.size) / 200), 2)
|
77 |
draw.rectangle(box, width=line_thickness, outline=tuple(color)) # plot
|
78 |
if label:
|
79 |
+
fontsize = max(round(max(im.size) / 40), 12)
|
80 |
font = ImageFont.truetype("Arial.ttf", fontsize)
|
81 |
txt_width, txt_height = font.getsize(label)
|
82 |
draw.rectangle([box[0], box[1] - txt_height + 4, box[0] + txt_width, box[1]], fill=tuple(color))
|
83 |
draw.text((box[0], box[1] - txt_height + 1), label, fill=(255, 255, 255), font=font)
|
84 |
+
return np.asarray(im)
|
85 |
|
86 |
|
87 |
def plot_wh_methods(): # from utils.plots import *; plot_wh_methods()
|