fix: 🐞 add optional dependencies and improve package management in models and late import for torch to fix ci (#1251)

Signed-off-by: Onuralp SEZER <[email protected]>

Author: onuralpszr

pyproject.toml

@@ -47,7 +47,6 @@ maintainers = [
 [project.scripts]
 sahi = "sahi.cli:app"
 
-
 [dependency-groups]
 dev = [
     "pytest>=7.2.2,<9.0.0",
@@ -78,6 +77,14 @@ build = [
 ]
 
 [project.optional-dependencies]
+yolov5 = ["yolov5>=6.0.0,<8.0.0"]
+yolo = ["ultralytics>=8.0.0"]
+torch = ["torch", "torchvision"]
+transformers = ["transformers>=4.49.0;python_version>='3.9'",]
+roboflow = ["inference>=0.51.5;python_version>='3.12'", "rfdetr>=1.1.0;python_version>='3.12'"]
+onnx = ["onnx;python_version>='3.10'", "onnxruntime;python_version>='3.10'"]
+all = ["yolov5>=6.0.0,<8.0.0", "ultralytics>=8.0.0", "torch", "torchvision", "transformers>=4.49.0;python_version>='3.9'", "inference>=0.51.5;python_version>='3.12'", "rfdetr>=1.1.0;python_version>='3.12'", "onnx;python_version>='3.10'", "onnxruntime;python_version>='3.10'"]
+
 mmdet = [
   # MMDetection dependencies - only for Python 3.11 and not on macOS ARM64
   "mmengine;python_version=='3.11' and (platform_system!='Darwin' or platform_machine!='arm64')",

sahi/models/ultralytics.py

@@ -4,7 +4,6 @@
 
 import cv2
 import numpy as np
-import torch
 
 from sahi.logger import logger
 from sahi.models.base import DetectionModel
@@ -72,6 +71,9 @@ def perform_inference(self, image: np.ndarray):
         """
 
         # Confirm model is loaded
+
+        import torch
+
         if self.model is None:
             raise ValueError("Model is not loaded, load it by calling .load_model()")
 

sahi/models/yolov5.py

@@ -8,16 +8,17 @@
 from sahi.models.base import DetectionModel
 from sahi.prediction import ObjectPrediction
 from sahi.utils.compatibility import fix_full_shape_list, fix_shift_amount_list
-from sahi.utils.import_utils import check_package_minimum_version, check_requirements
+from sahi.utils.import_utils import check_package_minimum_version
 
 
 class Yolov5DetectionModel(DetectionModel):
-    def check_dependencies(self) -> None:
-        check_requirements(["torch", "yolov5"])
+    def __init__(self, *args, **kwargs):
+        existing_packages = getattr(self, "required_packages", None) or []
+        self.required_packages = [*list(existing_packages), "yolov5", "torch"]
+        super().__init__(*args, **kwargs)
 
     def load_model(self):
         """Detection model is initialized and set to self.model."""
-
         import yolov5
 
         try:
@@ -71,13 +72,8 @@ def num_categories(self):
     @property
     def has_mask(self):
         """Returns if model output contains segmentation mask."""
-        import yolov5
-        from packaging import version
 
-        if version.parse(yolov5.__version__) < version.parse("6.2.0"):
-            return False
-        else:
-            return False  # fix when yolov5 supports segmentation models
+        return False  # fix when yolov5 supports segmentation models
 
     @property
     def category_names(self):

sahi/predict.py

@@ -3,23 +3,14 @@
 import os
 import time
 from collections.abc import Generator
-
-from PIL import Image
-
-from sahi.logger import logger
-from sahi.utils.import_utils import is_available
-
-# TODO: This does nothing for this module. The issue named here does not exist
-# https://github.com/obss/sahi/issues/526
-if is_available("torch"):
-    import torch  # noqa: F401
-
 from functools import cmp_to_key
 
 import numpy as np
+from PIL import Image
 from tqdm import tqdm
 
 from sahi.auto_model import AutoDetectionModel
+from sahi.logger import logger
 from sahi.models.base import DetectionModel
 from sahi.postprocess.combine import (
     GreedyNMMPostprocess,

sahi/utils/torch_utils.py

@@ -2,14 +2,18 @@
 
 import re
 from os import environ
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
-import torch
 from PIL.Image import Image
 
+if TYPE_CHECKING:
+    import torch
+
 
 def empty_cuda_cache() -> None:
+    import torch
+
     torch.cuda.empty_cache()
 
 
@@ -22,6 +26,8 @@ def to_float_tensor(img: np.ndarray | Image) -> torch.Tensor:
     Returns:
         torch.tensor
     """
+    import torch
+
     nparray: np.ndarray
     if isinstance(img, np.ndarray):
         nparray = img
@@ -54,6 +60,7 @@ def select_device(device: str | None = None) -> torch.device:
 
     Inspired by https://github.com/ultralytics/yolov5/blob/6371de8879e7ad7ec5283e8b95cc6dd85d6a5e72/utils/torch_utils.py#L107
     """
+    import torch
 
     if device == "cuda" or device is None:
         device = "cuda:0"