🎨 [Update] the dataset check, only check using set

yolo/tools/data_loader.py

@@ -23,7 +23,6 @@ from yolo.tools.data_augmentation import (
     VerticalFlip,
 )
 from yolo.tools.dataset_preparation import prepare_dataset
-from yolo.tools.drawer import draw_bboxes
 from yolo.utils.dataset_utils import (
     create_image_metadata,
     locate_label_paths,
@@ -204,7 +203,7 @@ def create_dataloader(data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: st
         return StreamDataLoader(data_cfg)
 
     if dataset_cfg.auto_download:
-        prepare_dataset(dataset_cfg)
+        prepare_dataset(dataset_cfg, task)
 
     return YoloDataLoader(data_cfg, dataset_cfg, task)
 

yolo/tools/dataset_preparation.py

@@ -52,7 +52,7 @@ def check_files(directory, expected_count=None):
     return len(files) == expected_count if expected_count is not None else bool(files)
 
 
-def prepare_dataset(dataset_cfg: DatasetConfig):
+def prepare_dataset(dataset_cfg: DatasetConfig, task: str):
     """
     Prepares dataset by downloading and unzipping if necessary.
     """
@@ -60,8 +60,8 @@ def prepare_dataset(dataset_cfg: DatasetConfig):
     for data_type, settings in dataset_cfg.auto_download.items():
         base_url = settings["base_url"]
         for dataset_type, dataset_args in settings.items():
-            if dataset_type == "base_url":
-                continue  # Skip the base_url entry
+            if dataset_type != "annotations" and dataset_cfg.get(task, task) != dataset_type:
+                continue
             file_name = f"{dataset_args.get('file_name', dataset_type)}.zip"
             url = f"{base_url}{file_name}"
             local_zip_path = os.path.join(data_dir, file_name)