@@ -133,9 +133,14 @@ if __name__ == '__main__': | |||
# model = custom(path_or_model='path/to/model.pt') # custom example | |||
# Verify inference | |||
import numpy as np | |||
from PIL import Image | |||
imgs = [Image.open(x) for x in Path('data/images').glob('*.jpg')] | |||
results = model(imgs) | |||
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 | |||
results = model(imgs) # batched inference | |||
results.print() | |||
results.save() |
@@ -254,7 +254,7 @@ 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) if isinstance(img, np.ndarray) else img # from np | |||
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 |