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'