Skip to content

Commit

Permalink
PyTorch Hub cv2 .save() .show() bug fix (ultralytics#2831)
Browse files Browse the repository at this point in the history
* 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.
  • Loading branch information
glenn-jocher authored Apr 18, 2021
1 parent 671cd5f commit f0e9571
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 18 deletions.
10 changes: 6 additions & 4 deletions hubconf.py
Original file line number Diff line number Diff line change
Expand Up @@ -124,13 +124,15 @@ def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True):
# model = custom(path_or_model='path/to/model.pt') # custom example

# Verify inference
import cv2
import numpy as np
from PIL import Image

imgs = [Image.open('data/images/bus.jpg'), # PIL
'data/images/zidane.jpg', # filename
'https://github.com/ultralytics/yolov5/raw/master/data/images/bus.jpg', # URI
np.zeros((640, 480, 3))] # numpy
imgs = ['data/images/zidane.jpg', # filename
'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg', # URI
cv2.imread('data/images/bus.jpg')[:, :, ::-1], # OpenCV
Image.open('data/images/bus.jpg'), # PIL
np.zeros((320, 640, 3))] # numpy

results = model(imgs) # batched inference
results.print()
Expand Down
4 changes: 2 additions & 2 deletions models/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -240,7 +240,7 @@ 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:
# filename: imgs = 'data/samples/zidane.jpg'
# filename: imgs = 'data/images/zidane.jpg'
# URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg'
# OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)
# PIL: = Image.open('image.jpg') # HWC x(640,1280,3)
Expand Down Expand Up @@ -271,7 +271,7 @@ def forward(self, imgs, size=640, augment=False, profile=False):
shape0.append(s) # image shape
g = (size / max(s)) # gain
shape1.append([y * g for y in s])
imgs[i] = im # update
imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update
shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)] # inference shape
x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad
x = np.stack(x, 0) if n > 1 else x[0][None] # stack
Expand Down
26 changes: 14 additions & 12 deletions utils/plots.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,32 +54,34 @@ def butter_lowpass(cutoff, fs, order):
return filtfilt(b, a, data) # forward-backward filter


def plot_one_box(x, img, color=None, label=None, line_thickness=3):
# Plots one bounding box on image img
tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness
def plot_one_box(x, im, color=None, label=None, line_thickness=3):
# Plots one bounding box on image 'im' using OpenCV
assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to plot_on_box() input image.'
tl = line_thickness or round(0.002 * (im.shape[0] + im.shape[1]) / 2) + 1 # line/font thickness
color = color or [random.randint(0, 255) for _ in range(3)]
c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
cv2.rectangle(im, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
if label:
tf = max(tl - 1, 1) # font thickness
t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled
cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
cv2.rectangle(im, c1, c2, color, -1, cv2.LINE_AA) # filled
cv2.putText(im, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)


def plot_one_box_PIL(box, img, color=None, label=None, line_thickness=None):
img = Image.fromarray(img)
draw = ImageDraw.Draw(img)
line_thickness = line_thickness or max(int(min(img.size) / 200), 2)
def plot_one_box_PIL(box, im, color=None, label=None, line_thickness=None):
# Plots one bounding box on image 'im' using PIL
im = Image.fromarray(im)
draw = ImageDraw.Draw(im)
line_thickness = line_thickness or max(int(min(im.size) / 200), 2)
draw.rectangle(box, width=line_thickness, outline=tuple(color)) # plot
if label:
fontsize = max(round(max(img.size) / 40), 12)
fontsize = max(round(max(im.size) / 40), 12)
font = ImageFont.truetype("Arial.ttf", fontsize)
txt_width, txt_height = font.getsize(label)
draw.rectangle([box[0], box[1] - txt_height + 4, box[0] + txt_width, box[1]], fill=tuple(color))
draw.text((box[0], box[1] - txt_height + 1), label, fill=(255, 255, 255), font=font)
return np.asarray(img)
return np.asarray(im)


def plot_wh_methods(): # from utils.plots import *; plot_wh_methods()
Expand Down

0 comments on commit f0e9571

Please sign in to comment.