From 5321d585946aea1a1363e81856c2eb00651a5429 Mon Sep 17 00:00:00 2001
From: Laughing <61612323+Laughing-q@users.noreply.github.com>
Date: Fri, 19 Apr 2024 09:14:27 +0800
Subject: [PATCH] Lightweight YOLO-World tests (#10133)

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
---
 tests/test_python.py | 27 ++++++++++++++++-----------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/tests/test_python.py b/tests/test_python.py
index 7f0dc7c3..67223c05 100644
--- a/tests/test_python.py
+++ b/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,