소스 검색

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

5.0
Yann Defretin GitHub 3 년 전
부모
커밋
5f42643a53
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1개의 변경된 파일4개의 추가작업 그리고 4개의 파일을 삭제
  1. +4
    -4
      models/common.py

+ 4
- 4
models/common.py 파일 보기

@@ -11,7 +11,7 @@ from PIL import Image, ImageDraw

from utils.datasets import letterbox
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
@@ -254,10 +254,10 @@ class Detections:
n = (pred[:, -1] == c).sum() # detections per class
str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string
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
# str += '%s %.2f, ' % (names[int(cls)], conf) # label
ImageDraw.Draw(img).rectangle(box, width=4, outline=colors[int(cls) % 10]) # plot
label = f'{self.names[int(cls)]} {conf:.2f}'
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:
print(str.rstrip(', '))
if show:

Loading…
취소
저장