ultralytics 8.0.40 TensorRT metadata and Results visualizer (#1014)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: Bogdan Gheorghe <112427971+bogdan-galileo@users.noreply.github.com>
Co-authored-by: Ayush Chaurasia <ayush.chaurarsia@gmail.com>
Co-authored-by: Jaap van de Loosdrecht <jaap@vdlmv.nl>
Co-authored-by: Noobtoss <96134731+Noobtoss@users.noreply.github.com>
Co-authored-by: nerdyespresso <106761627+nerdyespresso@users.noreply.github.com>
This commit is contained in:
Glenn Jocher
2023-02-17 20:06:06 +01:00
committed by GitHub
parent e799592718
commit 9047d737f4
40 changed files with 576 additions and 280 deletions

View File

@ -3,7 +3,7 @@
import subprocess
from pathlib import Path
from ultralytics.yolo.utils import ROOT, SETTINGS
from ultralytics.yolo.utils import LINUX, ROOT, SETTINGS
MODEL = Path(SETTINGS['weights_dir']) / 'yolov8n'
CFG = 'yolov8n'
@ -73,3 +73,8 @@ def test_export_segment_torchscript():
def test_export_classify_torchscript():
run(f'yolo export model={MODEL}-cls.pt format=torchscript')
def test_export_detect_edgetpu(enabled=False):
if enabled and LINUX:
run(f'yolo export model={MODEL}.pt format=edgetpu')

View File

@ -1,6 +1,5 @@
# Ultralytics YOLO 🚀, GPL-3.0 license
import platform
from pathlib import Path
import cv2
@ -10,12 +9,11 @@ from PIL import Image
from ultralytics import YOLO
from ultralytics.yolo.data.build import load_inference_source
from ultralytics.yolo.utils import ROOT, SETTINGS
from ultralytics.yolo.utils import LINUX, ROOT, SETTINGS
MODEL = Path(SETTINGS['weights_dir']) / 'yolov8n.pt'
CFG = 'yolov8n.yaml'
SOURCE = ROOT / 'assets/bus.jpg'
MACOS = platform.system() == 'Darwin' # macOS environment
def test_model_forward():
@ -87,24 +85,6 @@ def test_train_pretrained():
def test_export_torchscript():
"""
Format Argument Suffix CPU GPU
0 PyTorch - .pt True True
1 TorchScript torchscript .torchscript True True
2 ONNX onnx .onnx True True
3 OpenVINO openvino _openvino_model True False
4 TensorRT engine .engine False True
5 CoreML coreml .mlmodel True False
6 TensorFlow SavedModel saved_model _saved_model True True
7 TensorFlow GraphDef pb .pb True True
8 TensorFlow Lite tflite .tflite True False
9 TensorFlow Edge TPU edgetpu _edgetpu.tflite False False
10 TensorFlow.js tfjs _web_model False False
11 PaddlePaddle paddle _paddle_model True True
"""
from ultralytics.yolo.engine.exporter import export_formats
print(export_formats())
model = YOLO(MODEL)
f = model.export(format='torchscript')
YOLO(f)(SOURCE) # exported model inference
@ -124,9 +104,25 @@ def test_export_openvino():
def test_export_coreml(): # sourcery skip: move-assign
model = YOLO(MODEL)
f = model.export(format='coreml')
if MACOS:
YOLO(f)(SOURCE) # model prediction only supported on macOS
model.export(format='coreml')
# if MACOS:
# YOLO(f)(SOURCE) # model prediction only supported on macOS
def test_export_tflite(enabled=False):
# TF suffers from install conflicts on Windows and macOS
if enabled and LINUX:
model = YOLO(MODEL)
f = model.export(format='tflite')
YOLO(f)(SOURCE)
def test_export_pb(enabled=False):
# TF suffers from install conflicts on Windows and macOS
if enabled and LINUX:
model = YOLO(MODEL)
f = model.export(format='pb')
YOLO(f)(SOURCE)
def test_export_paddle(enabled=False):
@ -145,9 +141,8 @@ def test_workflow():
model = YOLO(MODEL)
model.train(data="coco8.yaml", epochs=1, imgsz=32)
model.val()
print(model.metrics)
model.predict(SOURCE)
model.export(format="onnx", opset=12) # export a model to ONNX format
model.export(format="onnx") # export a model to ONNX format
def test_predict_callback_and_setup():
@ -170,3 +165,13 @@ def test_predict_callback_and_setup():
print('test_callback', bs)
boxes = result.boxes # Boxes object for bbox outputs
print(boxes)
def test_result():
model = YOLO("yolov8n-seg.pt")
img = str(ROOT / "assets/bus.jpg")
res = model([img, img])
res[0].numpy()
res[0].cpu().numpy()
resimg = res[0].visualize(show_conf=False)
print(resimg)