✅ [Pass] test of autodownload deploy weight

README.md

@@ -12,8 +12,8 @@ Welcome to the official implementation of YOLOv7 and YOLOv9. This repository wil
 - This is the official YOLO model implementation with an MIT License.
 - For quick deployment: you can enter directly in the terminal:
 ```shell
-$pip install git+https://github.com/WongKinYiu/yolov9mit.git
-$yolo task=inference task.source=0 # source could be a single file, video, image folder, webcam ID
+pip install git+git@github.com:WongKinYiu/YOLO.git
+yolo task=inference task.source=0 # source could be a single file, video, image folder, webcam ID
 ```
 
 ## Introduction
@@ -23,8 +23,8 @@ $yolo task=inference task.source=0 # source could be a single file, video, image
 ## Installation
 To get started with YOLOv9, clone this repository and install the required dependencies:
 ```shell
-git clone [email protected]:WongKinYiu/yolov9mit.git
-cd yolov9mit
+git clone [email protected]:WongKinYiu/YOLO.git
+cd YOLO
 pip install -r requirements.txt
 ```
 
@@ -49,19 +49,20 @@ To train YOLOv9 on your dataset:
 1. Modify the configuration file `data/config.yaml` to point to your dataset.
 2. Run the training script:
 ```shell
-python lazy.py task=train task.batch_size=8 model=v9-c
+python yolo/lazy.py task=train task.data.batch_size=8 model=v9-c
 ```
 
 ### Transfer Learning
 To perform transfer learning with YOLOv9:
 ```shell
-python lazy.py task=train task.batch_size=8 model=v9-c task.data.dataset={dataset_config}
+python yolo/lazy.py task=train task.data.batch_size=8 model=v9-c dataset={dataset_config} device={cpu, mps, cuda}
 ```
 
 ### Inference
 To evaluate the model performance, use:
 ```shell
-python lazy.py weights=v9-c.pt # if cloned from GitHub
+python yolo/lazy.py task=inference weight=weights/v9-c-deploy.pt model=v9-c-deploy # use deploy weight
+python python yolo/lazy.py task=inference # if cloned from GitHub
 yolo task=inference task.data.source={Any} # if pip installed
 ```
 

tests/test_utils/test_loss.py

@@ -24,7 +24,7 @@ def cfg() -> Config:
 @pytest.fixture
 def model(cfg: Config):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = create_model(cfg.model)
+    model = create_model(cfg.model, weight_path=None)
     return model.to(device)
 
 

yolo/model/yolo.py

@@ -130,12 +130,11 @@ def create_model(model_cfg: ModelConfig, class_num: int = 80, weight_path: str =
     model = YOLO(model_cfg, class_num)
     logger.info("✅ Success load model")
     if weight_path:
-        if os.path.exists(weight_path):
-            model.model.load_state_dict(torch.load(weight_path))
-            logger.info("✅ Success load model weight")
-        else:
+        if not os.path.exists(weight_path):
             logger.info(f"🌐 Weight {weight_path} not found, try downloading")
             prepare_weight(weight_path=weight_path)
+        model.model.load_state_dict(torch.load(weight_path))
+        logger.info("✅ Success load model weight")
 
     log_model_structure(model.model)
     draw_model(model=model)

yolo/tools/loss_functions.py

@@ -13,6 +13,7 @@ from yolo.utils.bounding_box_utils import BoxMatcher, Vec2Box, calculate_iou
 class BCELoss(nn.Module):
     def __init__(self) -> None:
         super().__init__()
+        # TODO: Refactor the device, should be assign by config
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.bce = BCEWithLogitsLoss(pos_weight=torch.tensor([1.0], device=device), reduction="none")