Yann Defretin commited on
Commit
5f42643
1 Parent(s): db28ce6

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

Browse files
Files changed (1) hide show
  1. 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
- # str += '%s %.2f, ' % (names[int(cls)], conf) # label
260
- ImageDraw.Draw(img).rectangle(box, width=4, outline=colors[int(cls) % 10]) # plot
 
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: