Yann Defretin
commited on
Commit
•
5f42643
1
Parent(s):
db28ce6
Unified hub and detect.py box and labels plotting (#2243)
Browse files- models/common.py +4 -4
models/common.py
CHANGED
@@ -11,7 +11,7 @@ from PIL import Image, ImageDraw
|
|
11 |
|
12 |
from utils.datasets import letterbox
|
13 |
from utils.general import non_max_suppression, make_divisible, scale_coords, xyxy2xywh
|
14 |
-
from utils.plots import color_list
|
15 |
|
16 |
|
17 |
def autopad(k, p=None): # kernel, padding
|
@@ -254,10 +254,10 @@ class Detections:
|
|
254 |
n = (pred[:, -1] == c).sum() # detections per class
|
255 |
str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string
|
256 |
if show or save or render:
|
257 |
-
img = Image.fromarray(img.astype(np.uint8)) if isinstance(img, np.ndarray) else img # from np
|
258 |
for *box, conf, cls in pred: # xyxy, confidence, class
|
259 |
-
|
260 |
-
|
|
|
261 |
if pprint:
|
262 |
print(str.rstrip(', '))
|
263 |
if show:
|
|
|
11 |
|
12 |
from utils.datasets import letterbox
|
13 |
from utils.general import non_max_suppression, make_divisible, scale_coords, xyxy2xywh
|
14 |
+
from utils.plots import color_list, plot_one_box
|
15 |
|
16 |
|
17 |
def autopad(k, p=None): # kernel, padding
|
|
|
254 |
n = (pred[:, -1] == c).sum() # detections per class
|
255 |
str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string
|
256 |
if show or save or render:
|
|
|
257 |
for *box, conf, cls in pred: # xyxy, confidence, class
|
258 |
+
label = f'{self.names[int(cls)]} {conf:.2f}'
|
259 |
+
plot_one_box(box, img, label=label, color=colors[int(cls) % 10])
|
260 |
+
img = Image.fromarray(img.astype(np.uint8)) if isinstance(img, np.ndarray) else img # from np
|
261 |
if pprint:
|
262 |
print(str.rstrip(', '))
|
263 |
if show:
|