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:
@ -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')
|
||||
|
@ -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)
|
||||
|
Reference in New Issue
Block a user