🗃️ [Update] the path for saving models

yolo/config/general.yaml

@@ -11,4 +11,4 @@ lucky_number: 10
 use_wandb: False
 use_TensorBoard: False
 
-weight: v9-c.pt
+weight: weights/v9-c.pt

yolo/model/yolo.py

@@ -134,7 +134,7 @@ def create_model(cfg: Config) -> YOLO:
             logger.info("✅ Success load model weight")
         else:
             logger.info(f"🌐 Weight {cfg.weight} not found, try downloading")
-            prepare_weight(weight_name=cfg.weight)
+            prepare_weight(weight_path=cfg.weight)
 
     log_model_structure(model.model)
     draw_model(model=model)

yolo/tools/dataset_preparation.py

@@ -81,15 +81,19 @@ def prepare_dataset(cfg: DatasetConfig):
                 logger.error(f"Error verifying the {dataset_type} dataset after extraction.")
 
 
-def prepare_weight(downlaod_link: Optional[str] = None, weight_name: str = "v9-c.pt"):
+def prepare_weight(downlaod_link: Optional[str] = None, weight_path: str = "v9-c.pt"):
+    weight_name = os.path.basename(weight_path)
     if downlaod_link is None:
         downlaod_link = "https://github.com/WongKinYiu/yolov9mit/releases/download/v1.0-alpha/"
     weight_link = f"{downlaod_link}{weight_name}"
 
-    if os.path.exists(weight_name):
-        logger.info(f"Weight file '{weight_name}' already exists.")
+    if not os.path.isdir(os.path.dirname(weight_path)):
+        os.makedirs(os.path.dirname(weight_path))
+
+    if os.path.exists(weight_path):
+        logger.info(f"Weight file '{weight_path}' already exists.")
     try:
-        download_file(weight_link, weight_name)
+        download_file(weight_link, weight_path)
     except requests.exceptions.RequestException as e:
         logger.warning(f"Failed to download the weight file: {e}")