Add loguru & better path handling

README.md

@@ -3,7 +3,7 @@
 > \[!NOTE\]\
 > This is the latest version of tools CLI. If you are looking for the tools web application, please refer to the [web-app](https://github.com/luxonis/tools/tree/web-app) branch.
 
-This is a command-line tool that simplifies the conversion process of YOLO models. It supports the conversion of YOLOs ranging from V5 through V11 and Gold Yolo including oriented bounding boxes object detection (OBB), pose estimation, and instance segmentation variants of YOLOv8 and YOLO11 to ONNX format.
+This is a command-line tool that simplifies the conversion process of YOLO models. It supports the conversion of YOLOs ranging from V5 through V11 and Gold Yolo including oriented bounding boxes object detection (OBB), pose estimation, and instance segmentation variants of YOLOv8 and YOLO11 to ONNX format and archiving them in the NN Archive format.
 
 > \[!WARNING\]\
 > Please note that for the moment, we support conversion of YOLOv9 weights only from [Ultralytics](https://docs.ultralytics.com/models/yolov9/#performance-on-ms-coco-dataset).
@@ -34,8 +34,6 @@ cd tools
 
 ### Using Python package
 
-For this to work, you need to be located in the root folder of the project.
-
 ```bash
 # Install the package 
 pip install .

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tools"
-version = "0.2.1"
+version = "0.2.3"
 description = "Converter for YOLO models into .ONNX format."
 readme = "README.md"
 requires-python = ">=3.8"
@@ -49,7 +49,7 @@ select = ["E4", "E7", "E9", "F", "W", "B", "I", "FA"]
 convention = "google"
 
 [tool.mypy]
-python_version = "3.8"
+python_version = "3.10"
 ignore_missing_imports = true
 
 [tool.pyright]

requirements-dev.txt

@@ -1,3 +1,3 @@
-pre-commit==3.2.1
+pre-commit>=4.1.0
 pytest-cov>=4.1.0
 docker-squash>=1.1.0

requirements.txt

@@ -1,16 +1,16 @@
-torch>=1.9.1,<2.6.0
+torch>=2.0.0,<2.6.0
 torchvision>=0.10.1
 Pillow>=7.1.2
 PyYAML>=5.3.1
 gcsfs
-luxonis-ml[data,nn_archive]>=0.5.1
-# luxonis-ml[data,nn_archive] @ git+https://github.com/luxonis/luxonis-ml@dev
-onnx>=1.10.1
-onnxruntime
-onnxsim
+luxonis-ml[data,nn_archive,utils]>=0.6.1
+onnx>=1.17.0
+onnxruntime>=1.20.1
+onnxsim>=0.4.36
 s3fs
+tqdm
 s3transfer
 typer>=0.12.3
 psutil
-mmcv==1.5.0
-seaborn
+seaborn
+mmcv>=1.5.0,<2.0.0

tests/test_conversion.py

@@ -5,6 +5,7 @@
 import requests
 
 from tools.utils.version_detection import detect_version
+from tools.yolo.yolov6_exporter import YoloV6R4Exporter
 from tools.yolo.yolov8_exporter import YoloV8Exporter
 from tools.yolo.yolov10_exporter import YoloV10Exporter
 from tools.yolov7.yolov7_exporter import YoloV7Exporter
@@ -94,6 +95,14 @@ def test_yolov6nr4_automatic_version_detection():
     _remove_file("tests/yolov6n.pt")
 
 
+def test_yolov6nr4_model_conversion():
+    """Test the conversion of an YOLOv6nr4 model."""
+    _download_file(
+        "https://github.com/meituan/YOLOv6/releases/download/0.4.0/yolov6n.pt"
+    )
+    _test_model_conversion(YoloV6R4Exporter, "tests/yolov6n.pt", (640, 480), True)
+
+
 def test_yolov7t_automatic_version_detection():
     """Test the YOLOv7t autodetection of the model version."""
     _download_file(

tools/main.py

@@ -1,9 +1,11 @@
+#!/usr/bin/env python3
 from __future__ import annotations
 
-import logging
 from typing import Optional
 
 import typer
+from loguru import logger
+from luxonis_ml.utils import setup_logging
 from typing_extensions import Annotated
 
 from tools.utils import (
@@ -24,9 +26,7 @@
 )
 from tools.utils.constants import MISC_DIR
 
-logging.basicConfig()
-logging.getLogger().setLevel(logging.INFO)
-
+setup_logging()
 
 app = typer.Typer(help="Tools CLI", add_completion=False, rich_markup_mode="markdown")
 
@@ -80,9 +80,6 @@ def convert(
         ),
     ] = None,
 ):
-    logger = logging.getLogger(__name__)
-    logger.info("Converting model...")
-
     if version is not None and version not in YOLO_VERSIONS:
         logger.error("Wrong YOLO version selected!")
         raise typer.Exit(code=1) from None

tools/yolo/yolov10_exporter.py

@@ -1,12 +1,18 @@
 from __future__ import annotations
 
+import os
 import sys
 from typing import List, Optional, Tuple
 
+from loguru import logger
+
 from tools.modules import DetectV10, Exporter
 from tools.utils import get_first_conv2d_in_channels
 
-sys.path.append("./tools/yolo/ultralytics")
+current_dir = os.path.dirname(os.path.abspath(__file__))
+yolo_path = os.path.join(current_dir, "ultralytics")
+sys.path.append(yolo_path)
+
 from ultralytics.nn.modules import Detect  # noqa: E402
 from ultralytics.nn.tasks import attempt_load_one_weight  # noqa: E402
 
@@ -48,7 +54,7 @@ def load_model(self):
             self.number_of_channels = get_first_conv2d_in_channels(model)
             # print(f"Number of channels: {self.number_of_channels}")
         except Exception as e:
-            print(f"Error while getting number of channels: {e}")
+            logger.error(f"Error while getting number of channels: {e}")
 
         # check if image size is suitable
         gs = max(int(model.stride.max()), 32)  # model stride

tools/yolo/yolov5_exporter.py

@@ -1,18 +1,22 @@
 from __future__ import annotations
 
+import os
 import sys
 from typing import List, Optional, Tuple
 
 import torch.nn as nn
+from loguru import logger
 
 from tools.modules import DetectV5, Exporter
 from tools.utils import get_first_conv2d_in_channels
 
-sys.path.append("./tools/yolo/yolov5")
+current_dir = os.path.dirname(os.path.abspath(__file__))
+yolov5_path = os.path.join(current_dir, "yolov5")
+sys.path.append(yolov5_path)
 
 from models.common import Conv  # noqa: E402
-from models.experimental import attempt_load  # noqa: E402
-from models.yolo import Detect  # noqa: E402
+from models.experimental import attempt_load as attempt_load_yolov5  # noqa: E402
+from models.yolo import Detect as DetectYOLOv5  # noqa: E402
 from utils.activations import SiLU  # noqa: E402
 
 
@@ -35,7 +39,7 @@ def __init__(
     def load_model(self):
         # code based on export.py from YoloV5 repository
         # load the model
-        model = attempt_load(self.model_path, device="cpu")  # load FP32 model
+        model = attempt_load_yolov5(self.model_path, device="cpu")  # load FP32 model
 
         # check num classes and labels
         assert model.nc == len(
@@ -60,7 +64,7 @@ def load_model(self):
             if isinstance(m, Conv):  # assign export-friendly activations
                 if isinstance(m.act, nn.SiLU):
                     m.act = SiLU()
-            elif isinstance(m, Detect):
+            elif isinstance(m, DetectYOLOv5):
                 m.inplace = inplace
                 m.onnx_dynamic = False
                 if hasattr(m, "forward_export"):
@@ -78,7 +82,7 @@ def load_model(self):
             self.number_of_channels = get_first_conv2d_in_channels(model)
             # print(f"Number of channels: {self.number_of_channels}")
         except Exception as e:
-            print(f"Error while getting number of channels: {e}")
+            logger.error(f"Error while getting number of channels: {e}")
 
         self.m = model.module.model[-1] if hasattr(model, "module") else model.model[-1]
         self.num_branches = len(self.m.anchor_grid)

tools/yolo/yolov6_exporter.py

@@ -1,12 +1,18 @@
 from __future__ import annotations
 
+import os
 import sys
 from typing import Tuple
 
+from loguru import logger
+
 from tools.modules import DetectV6R4m, DetectV6R4s, Exporter
 from tools.utils import get_first_conv2d_in_channels
 
-sys.path.append("./tools/yolo/YOLOv6")
+current_dir = os.path.dirname(os.path.abspath(__file__))
+yolov6_path = os.path.join(current_dir, "YOLOv6")
+sys.path.append(yolov6_path)
+
 from yolov6.layers.common import RepVGGBlock  # noqa: E402
 from yolov6.models.heads.effidehead_distill_ns import Detect  # noqa: E402
 from yolov6.utils.checkpoint import load_checkpoint  # noqa: E402
@@ -51,7 +57,7 @@ def load_model(self):
             self.number_of_channels = get_first_conv2d_in_channels(model)
             # print(f"Number of channels: {self.number_of_channels}")
         except Exception as e:
-            print(f"Error while getting number of channels: {e}")
+            logger.error(f"Error while getting number of channels: {e}")
 
         self.num_branches = len(model.detect.grid)
 

tools/yolo/yolov8_exporter.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
+import os
 import sys
 from typing import List, Optional, Tuple
 
 import torch
+from loguru import logger
 from luxonis_ml.nn_archive import ArchiveGenerator
 from luxonis_ml.nn_archive.config_building_blocks import (
     DataType,
@@ -25,7 +27,10 @@
 )
 from tools.utils import get_first_conv2d_in_channels
 
-sys.path.append("./tools/yolo/ultralytics")
+current_dir = os.path.dirname(os.path.abspath(__file__))
+yolo_path = os.path.join(current_dir, "ultralytics")
+sys.path.append(yolo_path)
+
 from ultralytics.nn.modules import OBB, Classify, Detect, Pose, Segment  # noqa: E402
 from ultralytics.nn.tasks import attempt_load_one_weight  # noqa: E402
 
@@ -125,7 +130,7 @@ def load_model(self):
             self.number_of_channels = get_first_conv2d_in_channels(model)
             # print(f"Number of channels: {self.number_of_channels}")
         except Exception as e:
-            print(f"Error while getting number of channels: {e}")
+            logger.error(f"Error while getting number of channels: {e}")
 
         # Get output names
         self.all_output_names = get_output_names(self.mode)

tools/yolov6r1/yolov6_r1_exporter.py

@@ -1,12 +1,18 @@
 from __future__ import annotations
 
+import os
 import sys
 from typing import Tuple
 
+from loguru import logger
+
 from tools.modules import DetectV6R1, Exporter
 from tools.utils import get_first_conv2d_in_channels
 
-sys.path.append("./tools/yolov6r1/YOLOv6R1")
+current_dir = os.path.dirname(os.path.abspath(__file__))
+yolo_path = os.path.join(current_dir, "YOLOv6R1")
+sys.path.append(yolo_path)
+
 from yolov6.layers.common import RepVGGBlock  # noqa: E402
 from yolov6.utils.checkpoint import load_checkpoint  # noqa: E402
 
@@ -51,7 +57,7 @@ def load_model(self):
             self.number_of_channels = get_first_conv2d_in_channels(model)
             # print(f"Number of channels: {self.number_of_channels}")
         except Exception as e:
-            print(f"Error while getting number of channels: {e}")
+            logger.error(f"Error while getting number of channels: {e}")
 
         # check if image size is suitable
         gs = 2 ** (2 + self.num_branches)  # 1 = 8, 2 = 16, 3 = 32

tools/yolov6r3/gold_yolo_exporter.py

@@ -1,16 +1,23 @@
 from __future__ import annotations
 
+import os
 import sys
 from typing import Tuple
 
+from loguru import logger
+
 from tools.modules import DetectV6R3, Exporter
 from tools.utils import get_first_conv2d_in_channels
 
-sys.path.append("./tools/yolov6r3/Efficient-Computing/Detection/Gold-YOLO/")
-sys.path.append("./tools/yolov6r3/Efficient-Computing/Detection/Gold-YOLO/gold_yolo/")
+current_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.join(current_dir, "Efficient-Computing/Detection/Gold-YOLO/"))
+sys.path.append(
+    os.path.join(current_dir, "Efficient-Computing/Detection/Gold-YOLO/gold_yolo/")
+)
 sys.path.append(
-    "./tools/yolov6r3/Efficient-Computing/Detection/Gold-YOLO/yolov6/utils/"
+    os.path.join(current_dir, "Efficient-Computing/Detection/Gold-YOLO/yolov6/utils/")
 )
+
 from checkpoint import load_checkpoint as load_checkpoint_gold_yolo  # noqa: E402
 from switch_tool import switch_to_deploy  # noqa: E402
 
@@ -45,7 +52,7 @@ def load_model(self):
             self.number_of_channels = get_first_conv2d_in_channels(model)
             # print(f"Number of channels: {self.number_of_channels}")
         except Exception as e:
-            print(f"Error while getting number of channels: {e}")
+            logger.error(f"Error while getting number of channels: {e}")
 
         # check if image size is suitable
         gs = 2 ** (2 + self.num_branches)  # 1 = 8, 2 = 16, 3 = 32

tools/yolov6r3/yolov6_r3_exporter.py

@@ -1,12 +1,18 @@
 from __future__ import annotations
 
+import os
 import sys
 from typing import Tuple
 
+from loguru import logger
+
 from tools.modules import DetectV6R3, Exporter, YoloV6BackBone
 from tools.utils import get_first_conv2d_in_channels
 
-sys.path.append("tools/yolov6r3/YOLOv6R3")  # noqa: E402
+current_dir = os.path.dirname(os.path.abspath(__file__))
+yolo_path = os.path.join(current_dir, "YOLOv6R3")
+sys.path.append(yolo_path)  # noqa: E402
+
 from yolov6.layers.common import RepVGGBlock  # noqa: E402
 from yolov6.utils.checkpoint import load_checkpoint  # noqa: E402
 
@@ -74,7 +80,7 @@ def load_model(self):
             self.number_of_channels = get_first_conv2d_in_channels(model)
             # print(f"Number of channels: {self.number_of_channels}")
         except Exception as e:
-            print(f"Error while getting number of channels: {e}")
+            logger.error(f"Error while getting number of channels: {e}")
 
         # check if image size is suitable
         gs = 2 ** (2 + self.num_branches)  # 1 = 8, 2 = 16, 3 = 32