Add dota8.yaml and O tests (#7394)

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Laughing 2024-01-09 02:54:09 +08:00 committed by GitHub
parent d0562d7a2f
commit a6a2c256d4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
13 changed files with 176 additions and 16 deletions

View file

@ -13,12 +13,14 @@ 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)
('pose', 'yolov8n-pose', 'coco8-pose.yaml'),
('obb', 'yolov8n-obb', 'dota8.yaml'), ] # (task, model, data)
EXPORT_ARGS = [
('yolov8n', 'torchscript'),
('yolov8n-seg', 'torchscript'),
('yolov8n-cls', 'torchscript'),
('yolov8n-pose', 'torchscript'), ] # (model, format)
('yolov8n-pose', 'torchscript'),
('yolov8n-obb', 'torchscript'), ] # (model, format)
def run(cmd):

View file

@ -77,6 +77,7 @@ def test_predict_img():
seg_model = YOLO(WEIGHTS_DIR / 'yolov8n-seg.pt')
cls_model = YOLO(WEIGHTS_DIR / 'yolov8n-cls.pt')
pose_model = YOLO(WEIGHTS_DIR / 'yolov8n-pose.pt')
obb_model = YOLO(WEIGHTS_DIR / 'yolov8n-obb.pt')
im = cv2.imread(str(SOURCE))
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
@ -105,6 +106,8 @@ def test_predict_img():
assert len(results) == t.shape[0]
results = pose_model(t, imgsz=32)
assert len(results) == t.shape[0]
results = obb_model(t, imgsz=32)
assert len(results) == t.shape[0]
def test_predict_grey_and_4ch():