You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
128 lines
4.5 KiB
128 lines
4.5 KiB
# Ultralytics YOLO 🚀, AGPL-3.0 license |
|
|
|
from ultralytics import YOLO |
|
from ultralytics.cfg import get_cfg |
|
from ultralytics.engine.exporter import Exporter |
|
from ultralytics.models.yolo import classify, detect, segment |
|
from ultralytics.utils import ASSETS, DEFAULT_CFG, WEIGHTS_DIR |
|
|
|
CFG_DET = "yolov8n.yaml" |
|
CFG_SEG = "yolov8n-seg.yaml" |
|
CFG_CLS = "yolov8n-cls.yaml" # or 'squeezenet1_0' |
|
CFG = get_cfg(DEFAULT_CFG) |
|
MODEL = WEIGHTS_DIR / "yolov8n" |
|
|
|
|
|
def test_func(*args): # noqa |
|
"""Test function callback.""" |
|
print("callback test passed") |
|
|
|
|
|
def test_export(): |
|
"""Test model exporting functionality.""" |
|
exporter = Exporter() |
|
exporter.add_callback("on_export_start", test_func) |
|
assert test_func in exporter.callbacks["on_export_start"], "callback test failed" |
|
f = exporter(model=YOLO(CFG_DET).model) |
|
YOLO(f)(ASSETS) # exported model inference |
|
|
|
|
|
def test_detect(): |
|
"""Test object detection functionality.""" |
|
overrides = {"data": "coco8.yaml", "model": CFG_DET, "imgsz": 32, "epochs": 1, "save": False} |
|
CFG.data = "coco8.yaml" |
|
CFG.imgsz = 32 |
|
|
|
# Trainer |
|
trainer = detect.DetectionTrainer(overrides=overrides) |
|
trainer.add_callback("on_train_start", test_func) |
|
assert test_func in trainer.callbacks["on_train_start"], "callback test failed" |
|
trainer.train() |
|
|
|
# Validator |
|
val = detect.DetectionValidator(args=CFG) |
|
val.add_callback("on_val_start", test_func) |
|
assert test_func in val.callbacks["on_val_start"], "callback test failed" |
|
val(model=trainer.best) # validate best.pt |
|
|
|
# Predictor |
|
pred = detect.DetectionPredictor(overrides={"imgsz": [64, 64]}) |
|
pred.add_callback("on_predict_start", test_func) |
|
assert test_func in pred.callbacks["on_predict_start"], "callback test failed" |
|
result = pred(source=ASSETS, model=f"{MODEL}.pt") |
|
assert len(result), "predictor test failed" |
|
|
|
overrides["resume"] = trainer.last |
|
trainer = detect.DetectionTrainer(overrides=overrides) |
|
try: |
|
trainer.train() |
|
except Exception as e: |
|
print(f"Expected exception caught: {e}") |
|
return |
|
|
|
Exception("Resume test failed!") |
|
|
|
|
|
def test_segment(): |
|
"""Test image segmentation functionality.""" |
|
overrides = {"data": "coco8-seg.yaml", "model": CFG_SEG, "imgsz": 32, "epochs": 1, "save": False} |
|
CFG.data = "coco8-seg.yaml" |
|
CFG.imgsz = 32 |
|
# YOLO(CFG_SEG).train(**overrides) # works |
|
|
|
# Trainer |
|
trainer = segment.SegmentationTrainer(overrides=overrides) |
|
trainer.add_callback("on_train_start", test_func) |
|
assert test_func in trainer.callbacks["on_train_start"], "callback test failed" |
|
trainer.train() |
|
|
|
# Validator |
|
val = segment.SegmentationValidator(args=CFG) |
|
val.add_callback("on_val_start", test_func) |
|
assert test_func in val.callbacks["on_val_start"], "callback test failed" |
|
val(model=trainer.best) # validate best.pt |
|
|
|
# Predictor |
|
pred = segment.SegmentationPredictor(overrides={"imgsz": [64, 64]}) |
|
pred.add_callback("on_predict_start", test_func) |
|
assert test_func in pred.callbacks["on_predict_start"], "callback test failed" |
|
result = pred(source=ASSETS, model=f"{MODEL}-seg.pt") |
|
assert len(result), "predictor test failed" |
|
|
|
# Test resume |
|
overrides["resume"] = trainer.last |
|
trainer = segment.SegmentationTrainer(overrides=overrides) |
|
try: |
|
trainer.train() |
|
except Exception as e: |
|
print(f"Expected exception caught: {e}") |
|
return |
|
|
|
Exception("Resume test failed!") |
|
|
|
|
|
def test_classify(): |
|
"""Test image classification functionality.""" |
|
overrides = {"data": "imagenet10", "model": CFG_CLS, "imgsz": 32, "epochs": 1, "save": False} |
|
CFG.data = "imagenet10" |
|
CFG.imgsz = 32 |
|
# YOLO(CFG_SEG).train(**overrides) # works |
|
|
|
# Trainer |
|
trainer = classify.ClassificationTrainer(overrides=overrides) |
|
trainer.add_callback("on_train_start", test_func) |
|
assert test_func in trainer.callbacks["on_train_start"], "callback test failed" |
|
trainer.train() |
|
|
|
# Validator |
|
val = classify.ClassificationValidator(args=CFG) |
|
val.add_callback("on_val_start", test_func) |
|
assert test_func in val.callbacks["on_val_start"], "callback test failed" |
|
val(model=trainer.best) |
|
|
|
# Predictor |
|
pred = classify.ClassificationPredictor(overrides={"imgsz": [64, 64]}) |
|
pred.add_callback("on_predict_start", test_func) |
|
assert test_func in pred.callbacks["on_predict_start"], "callback test failed" |
|
result = pred(source=ASSETS, model=trainer.best) |
|
assert len(result), "predictor test failed"
|
|
|