🔀 [Merge] branch 'main' into SETUP

.gitignore

@@ -42,6 +42,7 @@ htmlcov/
 .coverage
 .coverage.*
 .cache
+*.cache
 nosetests.xml
 coverage.xml
 *.cover
@@ -140,3 +141,6 @@ runs
 
 # Ignore npm packages (if using frontend components)
 node_modules/
+
+# Not ignore image for demo
+!demo/images/inference/*

examples/example_inference.py

@@ -0,0 +1,35 @@
+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
+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)
+
+    tester = ModelTester(cfg, model, save_path, device)
+    tester.solve(dataloader)
+
+
+if __name__ == "__main__":
+    main()

examples/example_train.py

@@ -3,30 +3,28 @@ from pathlib import Path
 
 import hydra
 import torch
-from loguru import logger
 
 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.dataset_preparation import prepare_dataset
-from yolo.tools.trainer import ModelTrainer
+from yolo.tools.solver import 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()
-    save_path = validate_log_directory(cfg.hyper.general, cfg.name)
-    if cfg.download.auto:
-        prepare_dataset(cfg.download)
-
+    save_path = validate_log_directory(cfg, cfg.name)
     dataloader = create_dataloader(cfg)
     # TODO: get_device or rank, for DDP mode
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    trainer = ModelTrainer(cfg, save_path, device)
-    trainer.train(dataloader, cfg.hyper.train.epoch)
+    device = torch.device(cfg.device)
+    model = get_model(cfg).to(device)
+
+    trainer = ModelTrainer(cfg, model, save_path, device)
+    trainer.solve(dataloader, cfg.task.epoch)
 
 
 if __name__ == "__main__":

examples/lazy.py

@@ -0,0 +1,37 @@
+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()

tests/test_model/test_yolo.py

@@ -19,6 +19,7 @@ def test_build_model():
         cfg = compose(config_name=config_name)
 
         OmegaConf.set_struct(cfg.model, False)
+        cfg.weight = None
         model = YOLO(cfg.model, 80)
         assert len(model.model) == 38
 
@@ -26,6 +27,7 @@ def test_build_model():
 def test_get_model():
     with initialize(config_path=config_path, version_base=None):
         cfg = compose(config_name=config_name)
+        cfg.weight = None
         model = get_model(cfg)
         assert isinstance(model, YOLO)
 

tests/test_utils/test_dataaugment.py

@@ -39,7 +39,7 @@ def test_compose():
         return image, boxes
 
     compose = AugmentationComposer([mock_transform, mock_transform])
-    img = Image.new("RGB", (10, 10), color="blue")
+    img = Image.new("RGB", (640, 640), color="blue")
     boxes = torch.tensor([[0, 0.2, 0.2, 0.8, 0.8]])
 
     transformed_img, transformed_boxes = compose(img, boxes)

yolo/config/config.py

@@ -135,6 +135,8 @@ class Config:
     use_wandb: bool
     use_TensorBoard: bool
 
+    weight: Optional[str]
+
 
 @dataclass
 class YOLOLayer(nn.Module):

yolo/config/general.yaml

@@ -1,4 +1,4 @@
-deivce: [0]
+device: 0
 cpu_num: 16
 
 class_num: 80
@@ -9,4 +9,6 @@ exist_ok: True
 
 lucky_number: 10
 use_wandb: False
-use_TensorBoard: False
+use_TensorBoard: False
+
+weight: v9-c.pt

yolo/config/task/dataset/coco.yaml

@@ -1,4 +1,6 @@
 path: data/coco
+train: train2017
+
 
 auto_download:
   images:

yolo/config/task/dataset/demo.yaml

@@ -0,0 +1,3 @@
+path: demo
+
+auto_download:

yolo/config/task/inference.yaml

@@ -0,0 +1,11 @@
+task: inference
+defaults:
+  - dataset: demo
+data:
+  batch_size: 16
+  shuffle: False
+  pin_memory: True
+  data_augment: {}
+nms:
+  min_confidence: 0.75
+  min_iou: 0.5

yolo/model/yolo.py

@@ -1,10 +1,13 @@
+import os
 from typing import Any, Dict, List, Union
 
+import torch
 import torch.nn as nn
 from loguru import logger
 from omegaconf import ListConfig, OmegaConf
 
 from yolo.config.config import Config, Model, YOLOLayer
+from yolo.tools.dataset_preparation import prepare_weight
 from yolo.tools.drawer import draw_model
 from yolo.utils.logging_utils import log_model_structure
 from yolo.utils.module_utils import get_layer_map
@@ -125,6 +128,14 @@ def get_model(cfg: Config) -> YOLO:
     OmegaConf.set_struct(cfg.model, False)
     model = YOLO(cfg.model, cfg.class_num)
     logger.info("✅ Success load model")
+    if cfg.weight:
+        if os.path.exists(cfg.weight):
+            model.model.load_state_dict(torch.load(cfg.weight))
+            logger.info("✅ Success load model weight")
+        else:
+            logger.info(f"🌐 Weight {cfg.weight} not found, try downloading")
+            prepare_weight(weight_name=cfg.weight)
+
     log_model_structure(model.model)
     draw_model(model=model)
     return model

yolo/tools/dataset_preparation.py

@@ -1,10 +1,10 @@
 import os
 import zipfile
+from typing import Optional
 
 import requests
-from hydra import main
 from loguru import logger
-from tqdm import tqdm
+from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
 
 from yolo.config.config import DatasetConfig
 
@@ -13,18 +13,24 @@ def download_file(url, destination):
     """
     Downloads a file from the specified URL to the destination path with progress logging.
     """
-    logger.info(f"Downloading {os.path.basename(destination)}...")
     with requests.get(url, stream=True) as response:
         response.raise_for_status()
         total_size = int(response.headers.get("content-length", 0))
-        progress = tqdm(total=total_size, unit="iB", unit_scale=True, desc=os.path.basename(destination), leave=True)
-
-        with open(destination, "wb") as file:
-            for data in response.iter_content(chunk_size=1024 * 1024):  # 1 MB chunks
-                file.write(data)
-                progress.update(len(data))
-        progress.close()
-    logger.info("Download completed.")
+        with Progress(
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            "[progress.percentage]{task.percentage:>3.1f}%",
+            "•",
+            "{task.completed}/{task.total} bytes",
+            "•",
+            TimeRemainingColumn(),
+        ) as progress:
+            task = progress.add_task(f"📥 Downloading {os.path.basename(destination)}...", total=total_size)
+            with open(destination, "wb") as file:
+                for data in response.iter_content(chunk_size=1024 * 1024):  # 1 MB chunks
+                    file.write(data)
+                    progress.update(task, advance=len(data))
+    logger.info("✅ Download completed.")
 
 
 def unzip_file(source, destination):
@@ -46,7 +52,6 @@ def check_files(directory, expected_count=None):
     return len(files) == expected_count if expected_count is not None else bool(files)
 
 
-@main(config_path="../config/data", config_name="download", version_base=None)
 def prepare_dataset(cfg: DatasetConfig):
     """
     Prepares dataset by downloading and unzipping if necessary.
@@ -76,6 +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"):
+    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.")
+    try:
+        download_file(weight_link, weight_name)
+    except requests.exceptions.RequestException as e:
+        logger.warning(f"Failed to download the weight file: {e}")
+
+
 if __name__ == "__main__":
     import sys
 
@@ -83,4 +101,4 @@ if __name__ == "__main__":
     from utils.logging_utils import custom_logger
 
     custom_logger()
-    prepare_dataset()
+    prepare_weight()

yolo/tools/{trainer.py → solver.py}

@@ -6,8 +6,10 @@ from torch import Tensor
 from torch.cuda.amp import GradScaler, autocast
 
 from yolo.config.config import Config, TrainConfig
-from yolo.model.yolo import get_model
+from yolo.model.yolo import YOLO
+from yolo.tools.drawer import draw_bboxes
 from yolo.tools.loss_functions import get_loss_function
+from yolo.utils.bounding_box_utils import AnchorBoxConverter, bbox_nms
 from yolo.utils.logging_utils import ProgressTracker
 from yolo.utils.model_utils import (
     ExponentialMovingAverage,
@@ -17,16 +19,15 @@ from yolo.utils.model_utils import (
 
 
 class ModelTrainer:
-    def __init__(self, cfg: Config, save_path: str, device):
+    def __init__(self, cfg: Config, model: YOLO, save_path: str, device):
         train_cfg: TrainConfig = cfg.task
-        model = get_model(cfg)
-
-        self.model = model.to(device)
+        self.model = model
         self.device = device
         self.optimizer = create_optimizer(model, train_cfg.optimizer)
         self.scheduler = create_scheduler(self.optimizer, train_cfg.scheduler)
         self.loss_fn = get_loss_function(cfg)
-        self.progress = ProgressTracker(cfg, save_path, use_wandb=True)
+        self.progress = ProgressTracker(cfg, save_path, cfg.use_wandb)
+        self.num_epochs = cfg.task.epoch
 
         if getattr(train_cfg.ema, "enabled", False):
             self.ema = ExponentialMovingAverage(model, decay=train_cfg.ema.decay)
@@ -75,8 +76,9 @@ class ModelTrainer:
             self.ema.restore()
         torch.save(checkpoint, filename)
 
-    def train(self, dataloader, num_epochs):
+    def solve(self, dataloader):
         logger.info("🚄 Start Training!")
+        num_epochs = self.num_epochs
 
         with self.progress.progress:
             self.progress.start_train(num_epochs)
@@ -89,3 +91,27 @@ class ModelTrainer:
                 logger.info(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}")
                 if (epoch + 1) % 5 == 0:
                     self.save_checkpoint(epoch, f"checkpoint_epoch_{epoch+1}.pth")
+
+
+class ModelTester:
+    def __init__(self, cfg: Config, model: YOLO, save_path: str, device):
+        self.model = model
+        self.device = device
+        self.progress = ProgressTracker(cfg, save_path, cfg.use_wandb)
+
+        self.anchor2box = AnchorBoxConverter(cfg, device)
+        self.nms = cfg.task.nms
+        self.save_path = save_path
+
+    def solve(self, dataloader):
+        logger.info("👀 Start Inference!")
+
+        for images, _ in dataloader:
+            images = images.to(self.device)
+            with torch.no_grad():
+                raw_output = self.model(images)
+            predict, _ = self.anchor2box(raw_output[0][3:], with_logits=True)
+
+        nms_out = bbox_nms(predict, self.nms)
+        for image, bbox in zip(images, nms_out):
+            draw_bboxes(image, bbox, scaled_bbox=False, save_path=self.save_path)