diff --git a/hubconf.py b/hubconf.py index a2a43a7a10cb..d89502f4ee76 100644 --- a/hubconf.py +++ b/hubconf.py @@ -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() diff --git a/models/common.py b/models/common.py index 1130471e904b..2fdc0e0b70ca 100644 --- a/models/common.py +++ b/models/common.py @@ -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) @@ -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 diff --git a/utils/plots.py b/utils/plots.py index 5b23a34f5141..09b6bcd15a9f 100644 --- a/utils/plots.py +++ b/utils/plots.py @@ -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()