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:
@ -1,3 +1,5 @@
|
||||
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
|
@ -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
|
||||
|
@ -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
|
||||
|
@ -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))
|
||||
|
Reference in New Issue
Block a user