zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Unified hub and detect.py box and labels plotting (#2243)

This commit is contained in:
Yann Defretin 2021-02-19 21:38:05 +01:00 committed by GitHub
parent db28ce61ac
commit 5f42643a53
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
models/common.py
View File

@ -11,7 +11,7 @@ from PIL import Image, ImageDraw
from utils.datasets import letterbox from utils.datasets import letterbox
from utils.general import non_max_suppression, make_divisible, scale_coords, xyxy2xywh from utils.general import non_max_suppression, make_divisible, scale_coords, xyxy2xywh
from utils.plots import color_list from utils.plots import color_list, plot_one_box
def autopad(k, p=None): # kernel, padding def autopad(k, p=None): # kernel, padding
@ -254,10 +254,10 @@ class Detections:
n = (pred[:, -1] == c).sum() # detections per class n = (pred[:, -1] == c).sum() # detections per class
str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string
if show or save or render: if show or save or render:
img = Image.fromarray(img.astype(np.uint8)) if isinstance(img, np.ndarray) else img # from np
for *box, conf, cls in pred: # xyxy, confidence, class for *box, conf, cls in pred: # xyxy, confidence, class
# str += '%s %.2f, ' % (names[int(cls)], conf) # label label = f'{self.names[int(cls)]} {conf:.2f}'
ImageDraw.Draw(img).rectangle(box, width=4, outline=colors[int(cls) % 10]) # plot plot_one_box(box, img, label=label, color=colors[int(cls) % 10])
img = Image.fromarray(img.astype(np.uint8)) if isinstance(img, np.ndarray) else img # from np
if pprint: if pprint:
print(str.rstrip(', ')) print(str.rstrip(', '))
if show: if show: