Add Classification model YAML support (#154)

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>

tests/test_cli.py

@@ -3,8 +3,8 @@ from pathlib import Path
 
 from ultralytics.yolo.utils import ROOT, SETTINGS
 
-MODEL = Path(SETTINGS['weights_dir']) / 'yolov8n.pt'
-CFG = 'yolov8n.yaml'
+MODEL = Path(SETTINGS['weights_dir']) / 'yolov8n'
+CFG = 'yolov8n'
 
 
 def test_checks():
@@ -12,25 +12,25 @@ def test_checks():
 
 
 # Train checks ---------------------------------------------------------------------------------------------------------
-def test_train_detect():
-    os.system(f'yolo mode=train task=detect model={MODEL} data=coco128.yaml imgsz=32 epochs=1')
+def test_train_det():
+    os.system(f'yolo mode=train task=detect model={CFG}.yaml data=coco128.yaml imgsz=32 epochs=1')
 
 
-def test_train_segment():
-    os.system('yolo mode=train task=segment model=yolov8n-seg.yaml data=coco128-seg.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')
 
 
-def test_train_classify():
-    pass
+def test_train_cls():
+    os.system(f'yolo mode=train task=classify model={CFG}-cls.yaml data=imagenette160 imgsz=32 epochs=1')
 
 
 # Val checks -----------------------------------------------------------------------------------------------------------
 def test_val_detect():
-    os.system(f'yolo mode=val task=detect model={MODEL} data=coco128.yaml imgsz=32 epochs=1')
+    os.system(f'yolo mode=val task=detect model={MODEL}.pt data=coco128.yaml imgsz=32 epochs=1')
 
 
 def test_val_segment():
-    pass
+    os.system(f'yolo mode=val task=segment model={MODEL}-seg.pt data=coco128-seg.yaml imgsz=32 epochs=1')
 
 
 def test_val_classify():
@@ -39,11 +39,11 @@ def test_val_classify():
 
 # Predict checks -------------------------------------------------------------------------------------------------------
 def test_predict_detect():
-    os.system(f"yolo mode=predict model={MODEL} source={ROOT / 'assets'}")
+    os.system(f"yolo mode=predict model={MODEL}.pt source={ROOT / 'assets'}")
 
 
 def test_predict_segment():
-    pass
+    os.system(f"yolo mode=predict model={MODEL}-seg.pt source={ROOT / 'assets'}")
 
 
 def test_predict_classify():
@@ -52,11 +52,11 @@ def test_predict_classify():
 
 # Export checks --------------------------------------------------------------------------------------------------------
 def test_export_detect_torchscript():
-    os.system(f'yolo mode=export model={MODEL} format=torchscript')
+    os.system(f'yolo mode=export model={MODEL}.pt format=torchscript')
 
 
 def test_export_segment_torchscript():
-    pass
+    os.system(f'yolo mode=export model={MODEL}-seg.pt format=torchscript')
 
 
 def test_export_classify_torchscript():

tests/test_engine.py

@@ -71,7 +71,7 @@ def test_segment():
 def test_classify():
     overrides = {
         "data": "imagenette160",
-        "model": "squeezenet1_0",
+        "model": "yolov8n-cls.yaml",
         "imgsz": 32,
         "epochs": 1,
         "batch": 64,

tests/test_python.py

@@ -3,8 +3,8 @@ from pathlib import Path
 from ultralytics import YOLO
 from ultralytics.yolo.utils import ROOT, SETTINGS
 
-CFG = 'yolov8n.yaml'
 MODEL = Path(SETTINGS['weights_dir']) / 'yolov8n.pt'
+CFG = 'yolov8n.yaml'
 SOURCE = ROOT / 'assets/bus.jpg'