COCO8 and COCO8-seg Pytest and CI updates (#307)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: RangiLyu <lyuchqi@gmail.com>

tests/test_cli.py

@@ -15,11 +15,11 @@ def test_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')
+    os.system(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')
+    os.system(f'yolo mode=train task=segment model={CFG}-seg.yaml data=coco8-seg.yaml imgsz=32 epochs=1')
 
 
 def test_train_cls():
@@ -28,11 +28,11 @@ def test_train_cls():
 
 # Val checks -----------------------------------------------------------------------------------------------------------
 def test_val_detect():
-    os.system(f'yolo mode=val task=detect model={MODEL}.pt data=coco128.yaml imgsz=32 epochs=1')
+    os.system(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')
+    os.system(f'yolo mode=val task=segment model={MODEL}-seg.pt data=coco8-seg.yaml imgsz=32 epochs=1')
 
 
 def test_val_classify():

tests/test_engine.py

@@ -1,6 +1,5 @@
 # Ultralytics YOLO 🚀, GPL-3.0 license
 
-from ultralytics import YOLO
 from ultralytics.yolo.configs import get_config
 from ultralytics.yolo.utils import DEFAULT_CONFIG, ROOT
 from ultralytics.yolo.v8 import classify, detect, segment
@@ -13,9 +12,10 @@ SOURCE = ROOT / "assets"
 
 
 def test_detect():
-    overrides = {"data": "coco128.yaml", "model": CFG_DET, "imgsz": 32, "epochs": 1, "save": False}
-    CFG.data = "coco128.yaml"
-    # trainer
+    overrides = {"data": "coco8.yaml", "model": CFG_DET, "imgsz": 32, "epochs": 1, "save": False}
+    CFG.data = "coco8.yaml"
+
+    # Trainer
     trainer = detect.DetectionTrainer(overrides=overrides)
     trainer.train()
     trained_model = trainer.best
@@ -24,12 +24,10 @@ def test_detect():
     val = detect.DetectionValidator(args=CFG)
     val(model=trained_model)
 
-    # predictor
-    pred = detect.DetectionPredictor(overrides={"imgsz": [640, 640]})
-    i = 0
-    for _ in pred(source=SOURCE, model="yolov8n.pt", return_outputs=True):
-        i += 1
-    assert i == 2, "predictor test failed"
+    # Predictor
+    pred = detect.DetectionPredictor(overrides={"imgsz": [64, 64]})
+    result = pred(source=SOURCE, model="yolov8n.pt", return_outputs=True)
+    assert len(list(result)), "predictor test failed"
 
     overrides["resume"] = trainer.last
     trainer = detect.DetectionTrainer(overrides=overrides)
@@ -43,11 +41,11 @@ def test_detect():
 
 
 def test_segment():
-    overrides = {"data": "coco128-seg.yaml", "model": CFG_SEG, "imgsz": 32, "epochs": 1, "save": False}
-    CFG.data = "coco128-seg.yaml"
+    overrides = {"data": "coco8-seg.yaml", "model": CFG_SEG, "imgsz": 32, "epochs": 1, "save": False}
+    CFG.data = "coco8-seg.yaml"
     CFG.v5loader = False
+    # YOLO(CFG_SEG).train(**overrides)  # works
 
-    # YOLO(CFG_SEG).train(**overrides) # This works
     # trainer
     trainer = segment.SegmentationTrainer(overrides=overrides)
     trainer.train()
@@ -57,14 +55,12 @@ def test_segment():
     val = segment.SegmentationValidator(args=CFG)
     val(model=trained_model)
 
-    # predictor
-    pred = segment.SegmentationPredictor(overrides={"imgsz": [640, 640]})
-    i = 0
-    for _ in pred(source=SOURCE, model="yolov8n-seg.pt", return_outputs=True):
-        i += 1
-    assert i == 2, "predictor test failed"
+    # Predictor
+    pred = segment.SegmentationPredictor(overrides={"imgsz": [64, 64]})
+    result = pred(source=SOURCE, model="yolov8n-seg.pt", return_outputs=True)
+    assert len(list(result)) == 2, "predictor test failed"
 
-    # test resume
+    # Test resume
     overrides["resume"] = trainer.last
     trainer = segment.SegmentationTrainer(overrides=overrides)
     try:
@@ -81,8 +77,9 @@ def test_classify():
     CFG.data = "mnist160"
     CFG.imgsz = 32
     CFG.batch = 64
-    # YOLO(CFG_SEG).train(**overrides) # This works
-    # trainer
+    # YOLO(CFG_SEG).train(**overrides)  # works
+
+    # Trainer
     trainer = classify.ClassificationTrainer(overrides=overrides)
     trainer.train()
     trained_model = trainer.best
@@ -91,9 +88,7 @@ def test_classify():
     val = classify.ClassificationValidator(args=CFG)
     val(model=trained_model)
 
-    # predictor
-    pred = classify.ClassificationPredictor(overrides={"imgsz": [640, 640]})
-    i = 0
-    for _ in pred(source=SOURCE, model=trained_model, return_outputs=True):
-        i += 1
-    assert i == 2, "predictor test failed"
+    # Predictor
+    pred = classify.ClassificationPredictor(overrides={"imgsz": [64, 64]})
+    result = pred(source=SOURCE, model=trained_model, return_outputs=True)
+    assert len(list(result)) == 2, "predictor test failed"

tests/test_python.py

@@ -37,18 +37,18 @@ def test_predict_dir():
 
 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)
 
 
@@ -102,7 +102,7 @@ def test_all_model_yamls():
 
 def test_workflow():
     model = YOLO(MODEL)
-    model.train(data="coco128.yaml", epochs=1, imgsz=32)
+    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