New YOLOv8 Results() class for prediction outputs (#314)

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: Laughing-q <1185102784@qq.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: Viet Nhat Thai <60825385+vietnhatthai@users.noreply.github.com>
Co-authored-by: Paula Derrenger <107626595+pderrenger@users.noreply.github.com>

tests/test_cli.py

@@ -1,6 +1,6 @@
 # Ultralytics YOLO 🚀, GPL-3.0 license
 
-import os
+import subprocess
 from pathlib import Path
 
 from ultralytics.yolo.utils import ROOT, SETTINGS
@@ -9,30 +9,35 @@ MODEL = Path(SETTINGS['weights_dir']) / 'yolov8n'
 CFG = 'yolov8n'
 
 
+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=coco8.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=coco8-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=coco8.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=coco8-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 +46,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={ROOT / 'assets'}")
 
 
 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={ROOT / 'assets'}")
 
 
 def test_predict_classify():
@@ -54,12 +59,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_engine.py

@@ -28,8 +28,8 @@ def test_detect():
 
     # Predictor
     pred = detect.DetectionPredictor(overrides={"imgsz": [64, 64]})
-    result = pred(source=SOURCE, model=f"{MODEL}.pt", return_outputs=True)
-    assert len(list(result)), "predictor test failed"
+    result = pred(source=SOURCE, model=f"{MODEL}.pt")
+    assert len(result), "predictor test failed"
 
     overrides["resume"] = trainer.last
     trainer = detect.DetectionTrainer(overrides=overrides)
@@ -58,8 +58,8 @@ def test_segment():
 
     # Predictor
     pred = segment.SegmentationPredictor(overrides={"imgsz": [64, 64]})
-    result = pred(source=SOURCE, model=f"{MODEL}-seg.pt", return_outputs=True)
-    assert len(list(result)) == 2, "predictor test failed"
+    result = pred(source=SOURCE, model=f"{MODEL}-seg.pt")
+    assert len(result) == 2, "predictor test failed"
 
     # Test resume
     overrides["resume"] = trainer.last
@@ -90,5 +90,5 @@ def test_classify():
 
     # Predictor
     pred = classify.ClassificationPredictor(overrides={"imgsz": [64, 64]})
-    result = pred(source=SOURCE, model=trainer.best, return_outputs=True)
-    assert len(list(result)) == 2, "predictor test failed"
+    result = pred(source=SOURCE, model=trainer.best)
+    assert len(result) == 2, "predictor test failed"

tests/test_python.py

@@ -2,6 +2,10 @@
 
 from pathlib import Path
 
+import cv2
+import torch
+from PIL import Image
+
 from ultralytics import YOLO
 from ultralytics.yolo.utils import ROOT, SETTINGS
 
@@ -35,6 +39,21 @@ def test_predict_dir():
     model.predict(source=ROOT / "assets")
 
 
+def test_predict_img():
+    model = YOLO(MODEL)
+    img = Image.open(str(SOURCE))
+    output = model(source=img, save=True, verbose=True)  # PIL
+    assert len(output) == 1, "predict test failed"
+    img = cv2.imread(str(SOURCE))
+    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="coco8.yaml", imgsz=32)
@@ -106,3 +125,7 @@ def test_workflow():
     model.val()
     model.predict(SOURCE)
     model.export(format="onnx", opset=12)  # export a model to ONNX format
+
+
+if __name__ == "__main__":
+    test_predict_img()