Merge pull request #1514 from mikel-brostrom/refactor-reqs-checker

Refactor reqs checker
mikel-brostrom · Jul 12, 2024 · b0dc822 · b0dc822
2 parents 650b38d + b05bb5b
commit b0dc822
Show file tree

Hide file tree

Showing 13 changed files with 61 additions and 57 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -206,7 +206,7 @@ jobs:
         run: |
           python tracking/track.py --yolo-model yolov8n.pt --source $IMG --imgsz 320
           python tracking/track.py --yolo-model yolox_n.pt --source $IMG --imgsz 320
-          python tracking/track.py --yolo-model yolo_nas_s.pt --source $IMG --imgsz 320
+          # python tracking/track.py --yolo-model yolo_nas_s.pt --source $IMG --imgsz 320
 
 
   test-tracking-with-seg:

diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
@@ -32,4 +32,5 @@ jobs:
           push: true
           tags: |
             boxmot/boxmot:${{ env.VERSION }}
+            # will overwrite the previous image tagged as latest
             boxmot/boxmot:latest
diff --git a/boxmot/appearance/backends/onnx_backend.py b/boxmot/appearance/backends/onnx_backend.py
@@ -2,9 +2,9 @@
 from pathlib import Path
 
 from boxmot.appearance.backends.base_backend import BaseModelBackend
-from boxmot.utils.checks import TestRequirements
+from boxmot.utils.checks import RequirementsChecker
 
-tr = TestRequirements()
+checker = RequirementsChecker()
 
 
 class ONNXBackend(BaseModelBackend):
@@ -16,7 +16,7 @@ def __init__(self, weights, device, half):
 
     def load_model(self, w):
 
-        tr.check_packages(("onnxruntime-gpu==1.16.3" if self.cuda else "onnxruntime==1.16.3", ))
+        checker.check_packages(("onnxruntime-gpu==1.16.3" if self.cuda else "onnxruntime==1.16.3", ))
         import onnxruntime
 
         providers = (["CUDAExecutionProvider", "CPUExecutionProvider"] if self.cuda else ["CPUExecutionProvider"])

diff --git a/boxmot/appearance/backends/openvino_backend.py b/boxmot/appearance/backends/openvino_backend.py
@@ -3,9 +3,9 @@
 from boxmot.utils import logger as LOGGER
 
 from boxmot.appearance.backends.base_backend import BaseModelBackend
-from boxmot.utils.checks import TestRequirements
+from boxmot.utils.checks import RequirementsChecker
 
-tr = TestRequirements()
+checker = RequirementsChecker()
 
 
 class OpenVinoBackend(BaseModelBackend):
@@ -16,6 +16,7 @@ def __init__(self, weights, device, half):
         self.half = half
 
     def load_model(self, w):
+        checker.check_packages(("openvino-dev>=2022.3",))
 
         LOGGER.info(f"Loading {w} for OpenVINO inference...")
         try:

diff --git a/boxmot/appearance/backends/tensorrt_backend.py b/boxmot/appearance/backends/tensorrt_backend.py
@@ -4,9 +4,9 @@
 from boxmot.utils import logger as LOGGER
 
 from boxmot.appearance.backends.base_backend import BaseModelBackend
-from boxmot.utils.checks import TestRequirements
+from boxmot.utils.checks import RequirementsChecker
 
-tr = TestRequirements()
+checker = RequirementsChecker()
 
 
 class TensorRTBackend(BaseModelBackend):
@@ -19,7 +19,7 @@ def __init__(self, weights, device, half):
     def load_model(self, w):
 
         LOGGER.info(f"Loading {w} for TensorRT inference...")
-        tr.check_packages(("nvidia-tensorrt",))
+        checker.check_packages(("nvidia-tensorrt",))
         import tensorrt as trt  # https://developer.nvidia.com/nvidia-tensorrt-download
 
         if device.type == "cpu":

diff --git a/boxmot/appearance/backends/tflite_backend.py b/boxmot/appearance/backends/tflite_backend.py
@@ -4,9 +4,9 @@
 from boxmot.utils import logger as LOGGER
 
 from boxmot.appearance.backends.base_backend import BaseModelBackend
-from boxmot.utils.checks import TestRequirements
+from boxmot.utils.checks import RequirementsChecker
 
-tr = TestRequirements()
+checker = RequirementsChecker()
 
 
 class TFLiteBackend(BaseModelBackend):
@@ -17,11 +17,13 @@ def __init__(self, weights, device, half):
         self.half = half
 
     def load_model(self, w):
+        checker.check_packages(("tensorflow",))
 
         LOGGER.info(f"Loading {w} for TensorFlow Lite inference...")
-        import tensorflow as tf
-        interpreter = tf.lite.Interpreter(model_path=str(w))
+
         try:
+            import tensorflow as tf
+            interpreter = tf.lite.Interpreter(model_path=str(w))
             self.tf_lite_model = interpreter.get_signature_runner()
         except Exception as e:
             LOGGER.error(f'{e}. If SignatureDef error. Export you model with the official onn2tf docker')

diff --git a/boxmot/appearance/backends/torchscript_backend.py b/boxmot/appearance/backends/torchscript_backend.py
@@ -4,9 +4,6 @@
 from boxmot.utils import logger as LOGGER
 
 from boxmot.appearance.backends.base_backend import BaseModelBackend
-from boxmot.utils.checks import TestRequirements
-
-tr = TestRequirements()
 
 
 class TorchscriptBackend(BaseModelBackend):

diff --git a/boxmot/appearance/reid_export.py b/boxmot/appearance/reid_export.py
@@ -17,11 +17,11 @@
                                                  )
 from boxmot.utils import WEIGHTS
 from boxmot.utils import logger as LOGGER
-from boxmot.utils.checks import TestRequirements
+from boxmot.utils.checks import RequirementsChecker
 from boxmot.utils.torch_utils import select_device
 from boxmot.appearance.reid_auto_backend import ReidAutoBackend
 
-__tr = TestRequirements()
+checker = RequirementsChecker()
 
 
 def file_size(path):
@@ -56,7 +56,7 @@ def export_onnx(model, im, file, opset, dynamic, fp16, simplify):
     # ONNX export
     try:
         # required by onnx2tf
-        __tr.check_packages(("onnx==1.14.0",))
+        checker.check_packages(("onnx==1.14.0",))
         import onnx
 
         f = file.with_suffix(".onnx")
@@ -86,7 +86,7 @@ def export_onnx(model, im, file, opset, dynamic, fp16, simplify):
         if simplify:
             try:
                 cuda = torch.cuda.is_available()
-                __tr.check_packages(
+                checker.check_packages(
                     (
                         "onnxruntime-gpu" if cuda else "onnxruntime",
                         "onnx-simplifier>=0.4.1",
@@ -109,7 +109,7 @@ def export_onnx(model, im, file, opset, dynamic, fp16, simplify):
 
 
 def export_openvino(file, half):
-    __tr.check_packages(
+    checker.check_packages(
         ("openvino-dev>=2023.0",)
     )  # requires openvino-dev: https://pypi.org/project/openvino-dev/
     import openvino.runtime as ov  # noqa
@@ -136,7 +136,7 @@ def export_openvino(file, half):
 
 def export_tflite(file):
     try:
-        __tr.check_packages(
+        checker.check_packages(
             ("onnx2tf>=1.15.4", "tensorflow", "onnx_graphsurgeon>=0.3.26", "sng4onnx>=1.0.1"),
             cmds='--extra-index-url https://pypi.ngc.nvidia.com'
         )  # requires openvino-dev: https://pypi.org/project/openvino-dev/
@@ -162,7 +162,7 @@ def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose
             import tensorrt as trt
         except Exception:
             if platform.system() == "Linux":
-                __tr.check_packages(
+                checker.check_packages(
                     ["nvidia-tensorrt"],
                     cmds=("-U --index-url https://pypi.ngc.nvidia.com",),
                 )

diff --git a/boxmot/utils/checks.py b/boxmot/utils/checks.py
@@ -1,35 +1,38 @@
 # Mikel Broström 🔥 Yolo Tracking 🧾 AGPL-3.0 license
 
-import subprocess
-
 import pkg_resources
+import logging
+from pathlib import Path
+import pip
 
-from boxmot.utils import REQUIREMENTS, logger
-
+logger = logging.getLogger(__name__)
+REQUIREMENTS = Path('requirements.txt')
 
-class TestRequirements():
+class RequirementsChecker:
 
     def check_requirements(self):
         requirements = pkg_resources.parse_requirements(REQUIREMENTS.open())
         self.check_packages(requirements)
 
     def check_packages(self, requirements, cmds=''):
         """Test that each required package is available."""
-        # Ref: https://stackoverflow.com/a/45474387/
-
-        s = ''  # missing packages
+        missing_packages = []
         for r in requirements:
-            r = str(r)
             try:
-                pkg_resources.require(r)
+                pkg_resources.require(str(r))
             except Exception as e:
                 logger.error(f'{e}')
-                s += f'"{r}" '
-        if s:
-            logger.warning(f'\nMissing packages: {s}\nAtempting installation...')
-            try:
-                subprocess.check_output(f'pip install --no-cache {s} {cmds}', shell=True, stderr=subprocess.STDOUT)
-            except Exception as e:
-                logger.error(e)
-                exit()
-            logger.success('All the missing packages were installed successfully')
+                missing_packages.append(str(r))
+
+        if missing_packages:
+            self.install_packages(missing_packages, cmds)
+
+    def install_packages(self, packages, cmds=''):
+        try:
+            logger.warning(f'\nMissing packages: {", ".join(packages)}\nAttempting installation...')
+            pip_args = ['install', '--no-cache-dir'] + packages + cmds.split()
+            pip.main(pip_args)
+            logger.info('All the missing packages were installed successfully')
+        except Exception as e:
+            logger.error(f'Failed to install packages: {e}')
+            exit()
diff --git a/tracking/detectors/__init__.py b/tracking/detectors/__init__.py
@@ -1,9 +1,9 @@
 # Mikel Broström 🔥 Yolo Tracking 🧾 AGPL-3.0 license
 
 from boxmot.utils import logger as LOGGER
-from boxmot.utils.checks import TestRequirements
+from boxmot.utils.checks import RequirementsChecker
 
-tr = TestRequirements()
+checker = RequirementsChecker()
 
 
 def get_yolo_inferer(yolo_model):
@@ -13,9 +13,9 @@ def get_yolo_inferer(yolo_model):
             import yolox  # for linear_assignment
             assert yolox.__version__
         except (ImportError, AssertionError, AttributeError):
-            tr.check_packages(('yolox==0.3.0',), cmds='--no-dependencies')
-            tr.check_packages(('tabulate',))  # needed dependency
-            tr.check_packages(('thop',))  # needed dependency
+            checker.check_packages(('yolox==0.3.0',), cmds='--no-dependencies')
+            checker.check_packages(('tabulate',))  # needed dependency
+            checker.check_packages(('thop',))  # needed dependency
         from .yolox import YoloXStrategy
         return YoloXStrategy
     elif 'yolov8' in str(yolo_model):
@@ -27,7 +27,7 @@ def get_yolo_inferer(yolo_model):
             import super_gradients  # for linear_assignment
             assert super_gradients.__version__
         except (ImportError, AssertionError, AttributeError):
-            tr.check_packages(('super-gradients==3.1.3',))  # install
+            checker.check_packages(('super-gradients==3.1.3',))  # install
         from .yolonas import YoloNASStrategy
         return YoloNASStrategy
     else:

diff --git a/tracking/generate_dets_n_embs.py b/tracking/generate_dets_n_embs.py
@@ -8,12 +8,12 @@
 import torch
 
 from boxmot.utils import ROOT, WEIGHTS, logger as LOGGER
-from boxmot.utils.checks import TestRequirements
+from boxmot.utils.checks import RequirementsChecker
 from tracking.detectors import get_yolo_inferer
 from boxmot.appearance.reid_auto_backend import ReidAutoBackend
 
-__tr = TestRequirements()
-__tr.check_packages(('ultralytics @ git+https://github.com/mikel-brostrom/ultralytics.git', ))  # install
+checker = RequirementsChecker()
+checker.check_packages(('ultralytics @ git+https://github.com/mikel-brostrom/ultralytics.git', ))  # install
 
 from ultralytics import YOLO
 from ultralytics.data.utils import VID_FORMATS

diff --git a/tracking/generate_mot_results.py b/tracking/generate_mot_results.py
@@ -14,7 +14,7 @@
 
 from ultralytics.utils.files import increment_path 
 from boxmot.utils import ROOT, WEIGHTS, TRACKER_CONFIGS
-from boxmot.utils.checks import TestRequirements
+from boxmot.utils.checks import RequirementsChecker
 from boxmot.utils import logger as LOGGER
 
 from ultralytics.data.loaders import LoadImages
@@ -24,8 +24,8 @@
 from tracking.utils import convert_to_mot_format, write_mot_results
 from boxmot.utils.torch_utils import select_device
 
-__tr = TestRequirements()
-__tr.check_packages(('ultralytics @ git+https://github.com/mikel-brostrom/ultralytics.git', ))  # install
+checker = RequirementsChecker()
+checker.check_packages(('ultralytics @ git+https://github.com/mikel-brostrom/ultralytics.git', ))  # install
 
 
 def generate_mot_results(args):

diff --git a/tracking/track.py b/tracking/track.py
@@ -11,11 +11,11 @@
 from boxmot import TRACKERS
 from boxmot.tracker_zoo import create_tracker
 from boxmot.utils import ROOT, WEIGHTS, TRACKER_CONFIGS
-from boxmot.utils.checks import TestRequirements
+from boxmot.utils.checks import RequirementsChecker
 from tracking.detectors import get_yolo_inferer
 
-__tr = TestRequirements()
-__tr.check_packages(('ultralytics @ git+https://github.com/mikel-brostrom/ultralytics.git', ))  # install
+checker = RequirementsChecker()
+checker.check_packages(('ultralytics @ git+https://github.com/mikel-brostrom/ultralytics.git', ))  # install
 
 from ultralytics import YOLO
 from ultralytics.utils.plotting import Annotator, colors