Lightweight YOLO-World tests (#10133)

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
pull/10152/head^2
Laughing 7 months ago committed by GitHub
parent dae4ce8307
commit 5321d58594
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 27
      tests/test_python.py

@ -640,26 +640,31 @@ def test_yolo_world():
model.set_classes(["tree", "window"])
model(ASSETS / "bus.jpg", conf=0.01)
# Training from yaml
model = YOLO("yolov8s-worldv2.yaml") # no YOLOv8n-world model yet
model.train(data="coco8.yaml", epochs=2, imgsz=32, cache="disk", batch=-1, close_mosaic=1, name="yolo-world")
model = YOLO("yolov8s-worldv2.pt") # no YOLOv8n-world model yet
# val
model.val(data="coco8.yaml", imgsz=32, save_txt=True, save_json=True)
# Training from pretrain
model.train(data="coco8.yaml", epochs=2, imgsz=32, cache="disk", batch=-1, close_mosaic=1, name="yolo-world")
# Training from pretrain, evaluation process is included at the final stage of training.
# Use dota8.yaml which has less categories to reduce the inference time of CLIP model
model.train(
data="dota8.yaml",
epochs=1,
imgsz=32,
cache="disk",
batch=4,
close_mosaic=1,
name="yolo-world",
save_txt=True,
save_json=True,
)
# test WorWorldTrainerFromScratch
from ultralytics.models.yolo.world.train_world import WorldTrainerFromScratch
model = YOLO("yolov8s-worldv2.yaml") # no YOLOv8n-world model yet
model.train(
data={"train": {"yolo_data": ["coco8.yaml"]}, "val": {"yolo_data": ["coco8.yaml"]}},
epochs=2,
data={"train": {"yolo_data": ["dota8.yaml"]}, "val": {"yolo_data": ["dota8.yaml"]}},
epochs=1,
imgsz=32,
cache="disk",
batch=-1,
batch=4,
close_mosaic=1,
name="yolo-world",
trainer=WorldTrainerFromScratch,

Loading…
Cancel
Save