glenn-jocher commited on
Commit
c15e25c
·
unverified ·
1 Parent(s): aff03be

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.

Files changed (3) hide show
  1. hubconf.py +6 -4
  2. models/common.py +2 -2
  3. 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 = [Image.open('data/images/bus.jpg'), # PIL
131
- 'data/images/zidane.jpg', # filename
132
- 'https://github.com/ultralytics/yolov5/raw/master/data/images/bus.jpg', # URI
133
- np.zeros((640, 480, 3))] # numpy
 
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/samples/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,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, img, color=None, label=None, line_thickness=3):
58
- # Plots one bounding box on image img
59
- tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness
 
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(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
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(img, c1, c2, color, -1, cv2.LINE_AA) # filled
68
- cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
69
 
70
 
71
- def plot_one_box_PIL(box, img, color=None, label=None, line_thickness=None):
72
- img = Image.fromarray(img)
73
- draw = ImageDraw.Draw(img)
74
- line_thickness = line_thickness or max(int(min(img.size) / 200), 2)
 
75
  draw.rectangle(box, width=line_thickness, outline=tuple(color)) # plot
76
  if label:
77
- fontsize = max(round(max(img.size) / 40), 12)
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(img)
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()