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
GitHub
parent
d704507217
commit
9f6d48d3cf
10 changed files with 183 additions and 347 deletions
|
|
@ -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))
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue