Improve tests coverage and speed (#4340)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher
2023-08-13 22:24:01 +02:00
committed by GitHub
parent d704507217
commit 9f6d48d3cf
10 changed files with 183 additions and 347 deletions

View File

@ -1,3 +1,5 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license
import pytest

View File

@ -8,12 +8,16 @@ import pytest
from ultralytics.utils import ONLINE, ROOT, SETTINGS
WEIGHT_DIR = Path(SETTINGS['weights_dir'])
TASK_ARGS = [ # (task, model, data)
('detect', 'yolov8n', 'coco8.yaml'), ('segment', 'yolov8n-seg', 'coco8-seg.yaml'),
('classify', 'yolov8n-cls', 'imagenet10'), ('pose', 'yolov8n-pose', 'coco8-pose.yaml')]
EXPORT_ARGS = [ # (model, format)
('yolov8n', 'torchscript'), ('yolov8n-seg', 'torchscript'), ('yolov8n-cls', 'torchscript'),
('yolov8n-pose', 'torchscript')]
TASK_ARGS = [
('detect', 'yolov8n', 'coco8.yaml'),
('segment', 'yolov8n-seg', 'coco8-seg.yaml'),
('classify', 'yolov8n-cls', 'imagenet10'),
('pose', 'yolov8n-pose', 'coco8-pose.yaml'), ] # (task, model, data)
EXPORT_ARGS = [
('yolov8n', 'torchscript'),
('yolov8n-seg', 'torchscript'),
('yolov8n-cls', 'torchscript'),
('yolov8n-pose', 'torchscript'), ] # (model, format)
def run(cmd):
@ -22,9 +26,12 @@ def run(cmd):
def test_special_modes():
run('yolo checks')
run('yolo settings')
run('yolo help')
run('yolo checks')
run('yolo version')
run('yolo settings reset')
run('yolo copy-cfg')
run('yolo cfg')
@pytest.mark.parametrize('task,model,data', TASK_ARGS)
@ -34,21 +41,82 @@ def test_train(task, model, data):
@pytest.mark.parametrize('task,model,data', TASK_ARGS)
def test_val(task, model, data):
run(f'yolo val {task} model={model}.pt data={data} imgsz=32')
run(f'yolo val {task} model={WEIGHT_DIR / model}.pt data={data} imgsz=32')
@pytest.mark.parametrize('task,model,data', TASK_ARGS)
def test_predict(task, model, data):
run(f"yolo predict model={model}.pt source={ROOT / 'assets'} imgsz=32 save save_crop save_txt")
if ONLINE:
run(f'yolo predict model={model}.pt source=https://ultralytics.com/images/bus.jpg imgsz=32')
run(f'yolo predict model={model}.pt source=https://ultralytics.com/assets/decelera_landscape_min.mov imgsz=32')
run(f'yolo predict model={model}.pt source=https://ultralytics.com/assets/decelera_portrait_min.mov imgsz=32')
run(f"yolo predict model={WEIGHT_DIR / model}.pt source={ROOT / 'assets'} imgsz=32 save save_crop save_txt")
@pytest.mark.skipif(not ONLINE, reason='environment is offline')
@pytest.mark.parametrize('task,model,data', TASK_ARGS)
def test_predict_online(task, model, data):
mode = 'track' if task in ('detect', 'segment', 'pose') else 'predict' # mode for video inference
run(f'yolo predict model={WEIGHT_DIR / model}.pt source=https://ultralytics.com/images/bus.jpg imgsz=32')
run(f'yolo {mode} model={WEIGHT_DIR / model}.pt source=https://ultralytics.com/assets/decelera_landscape_min.mov imgsz=32'
)
# Run Python YouTube tracking because CLI is broken. TODO: fix CLI YouTube
# run(f'yolo {mode} model={model}.pt source=https://youtu.be/G17sBkb38XQ imgsz=32 tracker=bytetrack.yaml')
@pytest.mark.parametrize('model,format', EXPORT_ARGS)
def test_export(model, format):
run(f'yolo export model={model}.pt format={format}')
run(f'yolo export model={WEIGHT_DIR / model}.pt format={format} imgsz=32')
# Test SAM, RTDETR Models
def test_rtdetr(task='detect', model='yolov8n-rtdetr.yaml', data='coco8.yaml'):
# Warning: MUST use imgsz=640
run(f'yolo train {task} model={model} data={data} imgsz=640 epochs=1 cache=disk')
run(f'yolo val {task} model={model} data={data} imgsz=640')
run(f"yolo predict {task} model={model} source={ROOT / 'assets/bus.jpg'} imgsz=640 save save_crop save_txt")
def test_fastsam(task='segment', model='FastSAM-s.pt', data='coco8-seg.yaml'):
source = ROOT / 'assets/bus.jpg'
run(f'yolo segment val {task} model={model} data={data} imgsz=32')
run(f'yolo segment predict model={model} source={source} imgsz=32 save save_crop save_txt')
from ultralytics import FastSAM
from ultralytics.models.fastsam import FastSAMPrompt
# Create a FastSAM model
model = FastSAM('FastSAM-s.pt') # or FastSAM-x.pt
# Run inference on an image
everything_results = model(source, device='cpu', retina_masks=True, imgsz=1024, conf=0.4, iou=0.9)
# Everything prompt
prompt_process = FastSAMPrompt(source, everything_results, device='cpu')
ann = prompt_process.everything_prompt()
# Bbox default shape [0,0,0,0] -> [x1,y1,x2,y2]
ann = prompt_process.box_prompt(bbox=[200, 200, 300, 300])
# Text prompt
ann = prompt_process.text_prompt(text='a photo of a dog')
# Point prompt
# points default [[0,0]] [[x1,y1],[x2,y2]]
# point_label default [0] [1,0] 0:background, 1:foreground
ann = prompt_process.point_prompt(points=[[200, 200]], pointlabel=[1])
prompt_process.plot(annotations=ann, output='./')
def test_mobilesam():
from ultralytics import SAM
# Load the model
model = SAM('mobile_sam.pt')
# Predict a segment based on a point prompt
model.predict(ROOT / 'assets/zidane.jpg', points=[900, 370], labels=[1])
# Predict a segment based on a box prompt
model.predict(ROOT / 'assets/zidane.jpg', bboxes=[439, 437, 524, 709])
# Slow Tests

View File

@ -10,13 +10,13 @@ from ultralytics.utils import DEFAULT_CFG, ROOT, SETTINGS
CFG_DET = 'yolov8n.yaml'
CFG_SEG = 'yolov8n-seg.yaml'
CFG_CLS = 'squeezenet1_0'
CFG_CLS = 'yolov8n-cls.yaml' # or 'squeezenet1_0'
CFG = get_cfg(DEFAULT_CFG)
MODEL = Path(SETTINGS['weights_dir']) / 'yolov8n'
SOURCE = ROOT / 'assets'
def test_func(model=None):
def test_func(*args): # noqa
print('callback test passed')
@ -31,6 +31,7 @@ def test_export():
def test_detect():
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)
@ -65,6 +66,7 @@ def test_detect():
def test_segment():
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
@ -99,7 +101,7 @@ def test_segment():
def test_classify():
overrides = {'data': 'imagenet10', 'model': 'yolov8n-cls.yaml', 'imgsz': 32, 'epochs': 1, 'save': False}
overrides = {'data': 'imagenet10', 'model': CFG_CLS, 'imgsz': 32, 'epochs': 1, 'save': False}
CFG.data = 'imagenet10'
CFG.imgsz = 32
# YOLO(CFG_SEG).train(**overrides) # works

View File

@ -10,9 +10,11 @@ from torchvision.transforms import ToTensor
from ultralytics import RTDETR, YOLO
from ultralytics.data.build import load_inference_source
from ultralytics.utils import LINUX, ONLINE, ROOT, SETTINGS
from ultralytics.utils import LINUX, MACOS, ONLINE, ROOT, SETTINGS
from ultralytics.utils.torch_utils import TORCH_1_9
MODEL = Path(SETTINGS['weights_dir']) / 'path with spaces' / 'yolov8n.pt' # test spaces in path
WEIGHTS_DIR = Path(SETTINGS['weights_dir'])
MODEL = WEIGHTS_DIR / 'path with spaces' / 'yolov8n.pt' # test spaces in path
CFG = 'yolov8n.yaml'
SOURCE = ROOT / 'assets/bus.jpg'
SOURCE_GREYSCALE = Path(f'{SOURCE.parent / SOURCE.stem}_greyscale.jpg')
@ -26,39 +28,35 @@ im.convert('RGBA').save(SOURCE_RGBA) # 4-ch PNG with alpha
def test_model_forward():
model = YOLO(CFG)
model(SOURCE)
model(SOURCE, imgsz=32)
def test_model_info():
model = YOLO(CFG)
model.info()
model = YOLO(MODEL)
model.info(verbose=True)
def test_model_fuse():
model = YOLO(CFG)
model.fuse()
model = YOLO(MODEL)
model.fuse()
def test_predict_dir():
model = YOLO(MODEL)
model(source=ROOT / 'assets')
model(source=ROOT / 'assets', imgsz=32)
def test_predict_img():
model = YOLO(MODEL)
seg_model = YOLO('yolov8n-seg.pt')
cls_model = YOLO('yolov8n-cls.pt')
pose_model = YOLO('yolov8n-pose.pt')
seg_model = YOLO(WEIGHTS_DIR / 'yolov8n-seg.pt')
cls_model = YOLO(WEIGHTS_DIR / 'yolov8n-cls.pt')
pose_model = YOLO(WEIGHTS_DIR / 'yolov8n-pose.pt')
im = cv2.imread(str(SOURCE))
assert len(model(source=Image.open(SOURCE), save=True, verbose=True)) == 1 # PIL
assert len(model(source=im, save=True, save_txt=True)) == 1 # ndarray
assert len(model(source=[im, im], save=True, save_txt=True)) == 2 # batch
assert len(list(model(source=[im, im], save=True, stream=True))) == 2 # stream
assert len(model(torch.zeros(320, 640, 3).numpy())) == 1 # tensor to numpy
assert len(model(source=Image.open(SOURCE), save=True, verbose=True, imgsz=32)) == 1 # PIL
assert len(model(source=im, save=True, save_txt=True, imgsz=32)) == 1 # ndarray
assert len(model(source=[im, im], save=True, save_txt=True, imgsz=32)) == 2 # batch
assert len(list(model(source=[im, im], save=True, stream=True, imgsz=32))) == 2 # stream
assert len(model(torch.zeros(320, 640, 3).numpy(), imgsz=32)) == 1 # tensor to numpy
batch = [
str(SOURCE), # filename
Path(SOURCE), # Path
@ -66,20 +64,20 @@ def test_predict_img():
cv2.imread(str(SOURCE)), # OpenCV
Image.open(SOURCE), # PIL
np.zeros((320, 640, 3))] # numpy
assert len(model(batch, visualize=True)) == len(batch) # multiple sources in a batch
assert len(model(batch, imgsz=32)) == len(batch) # multiple sources in a batch
# Test tensor inference
im = cv2.imread(str(SOURCE)) # OpenCV
t = cv2.resize(im, (32, 32))
t = ToTensor()(t)
t = torch.stack([t, t, t, t])
results = model(t, visualize=True)
results = model(t, imgsz=32)
assert len(results) == t.shape[0]
results = seg_model(t, visualize=True)
results = seg_model(t, imgsz=32)
assert len(results) == t.shape[0]
results = cls_model(t, visualize=True)
results = cls_model(t, imgsz=32)
assert len(results) == t.shape[0]
results = pose_model(t, visualize=True)
results = pose_model(t, imgsz=32)
assert len(results) == t.shape[0]
@ -87,7 +85,13 @@ def test_predict_grey_and_4ch():
model = YOLO(MODEL)
for f in SOURCE_RGBA, SOURCE_GREYSCALE:
for source in Image.open(f), cv2.imread(str(f)), f:
model(source, save=True, verbose=True)
model(source, save=True, verbose=True, imgsz=32)
def test_track_stream():
# Test YouTube streaming inference (short 10 frame video) with non-default ByteTrack tracker
model = YOLO(MODEL)
model.track('https://youtu.be/G17sBkb38XQ', imgsz=32, tracker='bytetrack.yaml')
def test_val():
@ -95,11 +99,6 @@ def test_val():
model.val(data='coco8.yaml', imgsz=32)
def test_val_scratch():
model = YOLO(CFG)
model.val(data='coco8.yaml', imgsz=32)
def test_amp():
if torch.cuda.is_available():
from ultralytics.utils.checks import check_amp
@ -109,7 +108,7 @@ def test_amp():
def test_train_scratch():
model = YOLO(CFG)
model.train(data='coco8.yaml', epochs=1, imgsz=32, cache='disk') # test disk caching
model.train(data='coco8.yaml', epochs=1, imgsz=32, cache='disk', batch=-1) # test disk caching with AutoBatch
model(SOURCE)
@ -125,12 +124,6 @@ def test_export_torchscript():
YOLO(f)(SOURCE) # exported model inference
def test_export_torchscript_scratch():
model = YOLO(CFG)
f = model.export(format='torchscript')
YOLO(f)(SOURCE) # exported model inference
def test_export_onnx():
model = YOLO(MODEL)
f = model.export(format='onnx')
@ -138,14 +131,15 @@ def test_export_onnx():
def test_export_openvino():
model = YOLO(MODEL)
f = model.export(format='openvino')
YOLO(f)(SOURCE) # exported model inference
if not MACOS:
model = YOLO(MODEL)
f = model.export(format='openvino')
YOLO(f)(SOURCE) # exported model inference
def test_export_coreml(): # sourcery skip: move-assign
model = YOLO(MODEL)
model.export(format='coreml')
model.export(format='coreml', nms=True)
# if MACOS:
# YOLO(f)(SOURCE) # model prediction only supported on macOS
@ -174,9 +168,10 @@ def test_export_paddle(enabled=False):
def test_all_model_yamls():
for m in list((ROOT / 'models').rglob('yolo*.yaml')):
if m.name == 'yolov8-rtdetr.yaml': # except the rtdetr model
RTDETR(m.name)
for m in (ROOT / 'cfg' / 'models').rglob('*.yaml'):
if 'rtdetr' in m.name:
if TORCH_1_9: # torch<=1.8 issue - TypeError: __init__() got an unexpected keyword argument 'batch_first'
RTDETR(m.name)
else:
YOLO(m.name)
@ -190,10 +185,9 @@ def test_workflow():
def test_predict_callback_and_setup():
# test callback addition for prediction
# Test callback addition for prediction
def on_predict_batch_end(predictor): # results -> List[batch_size]
path, im0s, _, _ = predictor.batch
# print('on_predict_batch_end', im0s[0].shape)
im0s = im0s if isinstance(im0s, list) else [im0s]
bs = [predictor.dataset.bs for _ in range(len(path))]
predictor.results = zip(predictor.results, im0s, bs)
@ -204,42 +198,26 @@ def test_predict_callback_and_setup():
dataset = load_inference_source(source=SOURCE)
bs = dataset.bs # noqa access predictor properties
results = model.predict(dataset, stream=True) # source already setup
for _, (result, im0, bs) in enumerate(results):
for r, im0, bs in results:
print('test_callback', im0.shape)
print('test_callback', bs)
boxes = result.boxes # Boxes object for bbox outputs
boxes = r.boxes # Boxes object for bbox outputs
print(boxes)
def _test_results_api(res):
# General apis except plot
res = res.cpu().numpy()
# res = res.cuda()
res = res.to(device='cpu', dtype=torch.float32)
res.save_txt('label.txt', save_conf=False)
res.save_txt('label.txt', save_conf=True)
res.save_crop('crops/')
res.tojson(normalize=False)
res.tojson(normalize=True)
res.plot(pil=True)
res.plot(conf=True, boxes=False)
res.plot()
print(res)
print(res.path)
for k in res.keys:
print(getattr(res, k))
def test_results():
for m in ['yolov8n-pose.pt', 'yolov8n-seg.pt', 'yolov8n.pt', 'yolov8n-cls.pt']:
for m in 'yolov8n-pose.pt', 'yolov8n-seg.pt', 'yolov8n.pt', 'yolov8n-cls.pt':
model = YOLO(m)
res = model([SOURCE, SOURCE])
_test_results_api(res[0])
def test_track():
im = cv2.imread(str(SOURCE))
for m in ['yolov8n-pose.pt', 'yolov8n-seg.pt', 'yolov8n.pt']:
model = YOLO(m)
res = model.track(source=im)
_test_results_api(res[0])
results = model([SOURCE, SOURCE])
for r in results:
r = r.cpu().numpy()
r = r.to(device='cpu', dtype=torch.float32)
r.save_txt(txt_file='label.txt', save_conf=True)
r.save_crop(save_dir='crops/')
r.tojson(normalize=True)
r.plot(pil=True)
r.plot(conf=True, boxes=True)
print(r)
print(r.path)
for k in r.keys:
print(getattr(r, k))