🚸 [Add] auto download option for download dataset

config/config.py

@@ -8,6 +8,13 @@ class Model:
     model: Dict[str, List[Dict[str, Union[Dict, List, int]]]]
 
 
+@dataclass
+class Download:
+    auto: bool
+    path: str
+
+
 @dataclass
 class Config:
     model: Model
+    download: Download

config/data/download.yaml

@@ -1,3 +1,4 @@
+auto: True
 path: data/coco
 images:
   base_url: http://images.cocodataset.org/zips/

model/yolo.py

@@ -92,6 +92,7 @@ def get_model(model_cfg: dict) -> YOLO:
     Returns:
         YOLO: An instance of the model defined by the given configuration.
     """
+    OmegaConf.set_struct(model_cfg, False)
     model = YOLO(model_cfg)
     logger.info("✅ Success load model")
     return model

train.py

@@ -1,18 +1,17 @@
 from loguru import logger
 from model.yolo import get_model
 from tools.log_helper import custom_logger
-from utils.get_dataset import download_coco_dataset
+from utils.get_dataset import prepare_dataset
 import hydra
 from config.config import Config
-from omegaconf import OmegaConf
 
 
 @hydra.main(config_path="config", config_name="config", version_base=None)
 def main(cfg: Config):
-    OmegaConf.set_struct(cfg, False)
-    download_coco_dataset(cfg.download)
+    if cfg.download.auto:
+        prepare_dataset(cfg.download)
+
     model = get_model(cfg.model)
-    logger.info("Success load model")
 
 
 if __name__ == "__main__":

utils/get_dataset.py

@@ -44,7 +44,7 @@ def check_files(directory, expected_count):
 
 
 @hydra.main(config_path="../config/data", config_name="download", version_base=None)
-def download_coco_dataset(download_cfg):
+def prepare_dataset(download_cfg):
     data_dir = download_cfg.path
     base_url = download_cfg.images.base_url
     datasets = download_cfg.images.datasets