Spaces:
Sleeping
Sleeping
File size: 2,752 Bytes
0135475 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
# Ultralytics YOLO π, GPL-3.0 license
from pathlib import Path
from ultralytics.yolo.cfg import get_cfg
from ultralytics.yolo.utils import DEFAULT_CFG, ROOT, SETTINGS
from ultralytics.yolo.v8 import classify, detect, segment
CFG_DET = 'yolov8n.yaml'
CFG_SEG = 'yolov8n-seg.yaml'
CFG_CLS = 'squeezenet1_0'
CFG = get_cfg(DEFAULT_CFG)
MODEL = Path(SETTINGS['weights_dir']) / 'yolov8n'
SOURCE = ROOT / 'assets'
def test_detect():
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()
# Validator
val = detect.DetectionValidator(args=CFG)
val(model=trainer.best) # validate best.pt
# Predictor
pred = detect.DetectionPredictor(overrides={'imgsz': [64, 64]})
result = pred(source=SOURCE, 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():
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
# trainer
trainer = segment.SegmentationTrainer(overrides=overrides)
trainer.train()
# Validator
val = segment.SegmentationValidator(args=CFG)
val(model=trainer.best) # validate best.pt
# Predictor
pred = segment.SegmentationPredictor(overrides={'imgsz': [64, 64]})
result = pred(source=SOURCE, 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():
overrides = {'data': 'imagenet10', 'model': 'yolov8n-cls.yaml', '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.train()
# Validator
val = classify.ClassificationValidator(args=CFG)
val(model=trainer.best)
# Predictor
pred = classify.ClassificationPredictor(overrides={'imgsz': [64, 64]})
result = pred(source=SOURCE, model=trainer.best)
assert len(result), 'predictor test failed'
|