You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
138 lines
5.0 KiB
138 lines
5.0 KiB
# Ultralytics YOLO 🚀, AGPL-3.0 license |
|
|
|
import subprocess |
|
|
|
import pytest |
|
|
|
from ultralytics.utils import ASSETS, WEIGHTS_DIR |
|
from ultralytics.utils.checks import cuda_device_count, cuda_is_available |
|
|
|
CUDA_IS_AVAILABLE = cuda_is_available() |
|
CUDA_DEVICE_COUNT = cuda_device_count() |
|
TASK_ARGS = [ |
|
("detect", "yolov8n", "coco8.yaml"), |
|
("segment", "yolov8n-seg", "coco8-seg.yaml"), |
|
("classify", "yolov8n-cls", "imagenet10"), |
|
("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"), |
|
("yolov8n-obb", "torchscript"), |
|
] # (model, format) |
|
|
|
|
|
def run(cmd): |
|
"""Execute a shell command using subprocess.""" |
|
subprocess.run(cmd.split(), check=True) |
|
|
|
|
|
def test_special_modes(): |
|
"""Test various special command modes of YOLO.""" |
|
run("yolo help") |
|
run("yolo checks") |
|
run("yolo version") |
|
run("yolo settings reset") |
|
run("yolo cfg") |
|
|
|
|
|
@pytest.mark.parametrize("task,model,data", TASK_ARGS) |
|
def test_train(task, model, data): |
|
"""Test YOLO training for a given task, model, and data.""" |
|
run(f"yolo train {task} model={model}.yaml data={data} imgsz=32 epochs=1 cache=disk") |
|
|
|
|
|
@pytest.mark.parametrize("task,model,data", TASK_ARGS) |
|
def test_val(task, model, data): |
|
"""Test YOLO validation for a given task, model, and data.""" |
|
run(f"yolo val {task} model={WEIGHTS_DIR / model}.pt data={data} imgsz=32 save_txt save_json") |
|
|
|
|
|
@pytest.mark.parametrize("task,model,data", TASK_ARGS) |
|
def test_predict(task, model, data): |
|
"""Test YOLO prediction on sample assets for a given task and model.""" |
|
run(f"yolo predict model={WEIGHTS_DIR / model}.pt source={ASSETS} imgsz=32 save save_crop save_txt") |
|
|
|
|
|
@pytest.mark.parametrize("model,format", EXPORT_ARGS) |
|
def test_export(model, format): |
|
"""Test exporting a YOLO model to different formats.""" |
|
run(f"yolo export model={WEIGHTS_DIR / model}.pt format={format} imgsz=32") |
|
|
|
|
|
def test_rtdetr(task="detect", model="yolov8n-rtdetr.yaml", data="coco8.yaml"): |
|
"""Test the RTDETR functionality with the Ultralytics framework.""" |
|
# Warning: MUST use imgsz=640 |
|
run(f"yolo train {task} model={model} data={data} --imgsz= 640 epochs =1, cache = disk") # add coma, spaces to args |
|
run(f"yolo predict {task} model={model} source={ASSETS / 'bus.jpg'} imgsz=640 save save_crop save_txt") |
|
|
|
|
|
def test_fastsam(task="segment", model=WEIGHTS_DIR / "FastSAM-s.pt", data="coco8-seg.yaml"): |
|
"""Test FastSAM segmentation functionality within Ultralytics.""" |
|
source = 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 |
|
from ultralytics.models.sam import Predictor |
|
|
|
# Create a FastSAM model |
|
sam_model = FastSAM(model) # or FastSAM-x.pt |
|
|
|
# Run inference on an image |
|
everything_results = sam_model(source, device="cpu", retina_masks=True, imgsz=1024, conf=0.4, iou=0.9) |
|
|
|
# Remove small regions |
|
new_masks, _ = Predictor.remove_small_regions(everything_results[0].masks.data, min_area=20) |
|
|
|
# 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(): |
|
"""Test MobileSAM segmentation functionality using Ultralytics.""" |
|
from ultralytics import SAM |
|
|
|
# Load the model |
|
model = SAM(WEIGHTS_DIR / "mobile_sam.pt") |
|
|
|
# Source |
|
source = ASSETS / "zidane.jpg" |
|
|
|
# Predict a segment based on a point prompt |
|
model.predict(source, points=[900, 370], labels=[1]) |
|
|
|
# Predict a segment based on a box prompt |
|
model.predict(source, bboxes=[439, 437, 524, 709]) |
|
|
|
# Predict all |
|
# model(source) |
|
|
|
|
|
# Slow Tests ----------------------------------------------------------------------------------------------------------- |
|
@pytest.mark.slow |
|
@pytest.mark.parametrize("task,model,data", TASK_ARGS) |
|
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available") |
|
@pytest.mark.skipif(CUDA_DEVICE_COUNT < 2, reason="DDP is not available") |
|
def test_train_gpu(task, model, data): |
|
"""Test YOLO training on GPU(s) for various tasks and models.""" |
|
run(f"yolo train {task} model={model}.yaml data={data} imgsz=32 epochs=1 device=0") # single GPU |
|
run(f"yolo train {task} model={model}.pt data={data} imgsz=32 epochs=1 device=0,1") # multi GPU
|
|
|