🚚 [Rename] get_model to create, rename examples

examples/lazy.py

@@ -1,37 +0,0 @@
-import sys
-from pathlib import Path
-
-import hydra
-import torch
-
-project_root = Path(__file__).resolve().parent.parent
-sys.path.append(str(project_root))
-
-from yolo.config.config import Config
-from yolo.model.yolo import get_model
-from yolo.tools.data_loader import create_dataloader
-from yolo.tools.solver import ModelTester, ModelTrainer
-from yolo.utils.logging_utils import custom_logger, validate_log_directory
-
-
-@hydra.main(config_path="../yolo/config", config_name="config", version_base=None)
-def main(cfg: Config):
-    custom_logger()
-
-    custom_logger()
-    save_path = validate_log_directory(cfg, cfg.name)
-    dataloader = create_dataloader(cfg)
-    device = torch.device(cfg.device)
-    model = get_model(cfg).to(device)
-
-    if cfg.task.task == "train":
-        trainer = ModelTrainer(cfg, model, save_path, device)
-        trainer.solve(dataloader)
-
-    if cfg.task.task == "inference":
-        tester = ModelTester(cfg, model, save_path, device)
-        tester.solve(dataloader)
-
-
-if __name__ == "__main__":
-    main()

examples/notebook_inference.ipynb

@@ -0,0 +1,25 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

examples/notebook_train.ipynb

@@ -0,0 +1,25 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

examples/{example_inference.py → sample_inference.py}

@@ -8,7 +8,7 @@ project_root = Path(__file__).resolve().parent.parent
 sys.path.append(str(project_root))
 
 from yolo.config.config import Config
-from yolo.model.yolo import get_model
+from yolo.model.yolo import create_model
 from yolo.tools.data_loader import create_dataloader
 from yolo.tools.solver import ModelTester
 from yolo.utils.logging_utils import custom_logger, validate_log_directory
@@ -17,15 +17,11 @@ from yolo.utils.logging_utils import custom_logger, validate_log_directory
 @hydra.main(config_path="../yolo/config", config_name="config", version_base=None)
 def main(cfg: Config):
     custom_logger()
-    save_path = validate_log_directory(cfg, cfg.name)
-
-    device = torch.device(cfg.device)
-    model = get_model(cfg).to(device)
-
     save_path = validate_log_directory(cfg, cfg.name)
     dataloader = create_dataloader(cfg)
+
     device = torch.device(cfg.device)
-    model = get_model(cfg).to(device)
+    model = create_model(cfg).to(device)
 
     tester = ModelTester(cfg, model, save_path, device)
     tester.solve(dataloader)

examples/{example_train.py → sample_train.py}

@@ -8,7 +8,7 @@ project_root = Path(__file__).resolve().parent.parent
 sys.path.append(str(project_root))
 
 from yolo.config.config import Config
-from yolo.model.yolo import get_model
+from yolo.model.yolo import create_model
 from yolo.tools.data_loader import create_dataloader
 from yolo.tools.solver import ModelTrainer
 from yolo.utils.logging_utils import custom_logger, validate_log_directory
@@ -21,7 +21,7 @@ def main(cfg: Config):
     dataloader = create_dataloader(cfg)
     # TODO: get_device or rank, for DDP mode
     device = torch.device(cfg.device)
-    model = get_model(cfg).to(device)
+    model = create_model(cfg).to(device)
 
     trainer = ModelTrainer(cfg, model, save_path, device)
     trainer.solve(dataloader, cfg.task.epoch)

tests/test_model/test_yolo.py

@@ -8,7 +8,7 @@ from omegaconf import OmegaConf
 project_root = Path(__file__).resolve().parent.parent.parent
 sys.path.append(str(project_root))
 
-from yolo.model.yolo import YOLO, get_model
+from yolo.model.yolo import YOLO, create_model
 
 config_path = "../../yolo/config"
 config_name = "config"
@@ -24,18 +24,18 @@ def test_build_model():
         assert len(model.model) == 38
 
 
-def test_get_model():
+def test_create_model():
     with initialize(config_path=config_path, version_base=None):
         cfg = compose(config_name=config_name)
         cfg.weight = None
-        model = get_model(cfg)
+        model = create_model(cfg)
         assert isinstance(model, YOLO)
 
 
 def test_yolo_forward_output_shape():
     with initialize(config_path=config_path, version_base=None):
         cfg = compose(config_name=config_name)
-        model = get_model(cfg)
+        model = create_model(cfg)
         # 2 - batch size, 3 - number of channels, 640x640 - image dimensions
         dummy_input = torch.rand(2, 3, 640, 640)
 

yolo/lazy.py

@@ -8,7 +8,7 @@ project_root = Path(__file__).resolve().parent.parent
 sys.path.append(str(project_root))
 
 from yolo.config.config import Config
-from yolo.model.yolo import get_model
+from yolo.model.yolo import create_model
 from yolo.tools.data_loader import create_dataloader
 from yolo.tools.solver import ModelTester, ModelTrainer
 from yolo.utils.logging_utils import custom_logger, validate_log_directory
@@ -20,7 +20,7 @@ def main(cfg: Config):
     save_path = validate_log_directory(cfg, cfg.name)
     dataloader = create_dataloader(cfg)
     device = torch.device(cfg.device)
-    model = get_model(cfg).to(device)
+    model = create_model(cfg).to(device)
 
     if cfg.task.task == "train":
         trainer = ModelTrainer(cfg, model, save_path, device)

yolo/model/yolo.py

@@ -116,7 +116,7 @@ class YOLO(nn.Module):
             raise ValueError(f"Unsupported layer type: {layer_type}")
 
 
-def get_model(cfg: Config) -> YOLO:
+def create_model(cfg: Config) -> YOLO:
     """Constructs and returns a model from a Dictionary configuration file.
 
     Args:

yolo/tools/drawer.py

@@ -14,6 +14,7 @@ def draw_bboxes(
     *,
     scaled_bbox: bool = True,
     save_path: str = "",
+    save_name: str = "visualize.png",
 ):
     """
     Draw bounding boxes on an image.
@@ -46,7 +47,7 @@ def draw_bboxes(
         draw.rectangle(shape, outline="red", width=3)
         draw.text((x_min, y_min), str(int(class_id)), font=font, fill="blue")
 
-    save_image_path = os.path.join(save_path, "visualize.png")
+    save_image_path = os.path.join(save_path, save_name)
     img.save(save_image_path)  # Save the image with annotations
     logger.info(f"💾 Saved visualize image at {save_image_path}")
     return img
@@ -56,9 +57,9 @@ def draw_model(*, model_cfg=None, model=None, v7_base=False):
     from graphviz import Digraph
 
     if model_cfg:
-        from yolo.model.yolo import get_model
+        from yolo.model.yolo import create_model
 
-        model = get_model(model_cfg)
+        model = create_model(model_cfg)
     elif model is None:
         raise ValueError("Drawing Object is None")