85 enable gpu support for yoloonnxdetector

CHANGELOG.md

@@ -5,6 +5,21 @@
 - Added TextEmbeddings transformer, for compute embedding using SentenceTransformers
 
 
+## [0.2.6] - 19.11.2025
+
+### 🚀 Features
+
+- Enable support GPU in [YoloOnnxDetector](https://scaledp.stabrise.com/en/latest/models/detectors/yolo_onnx_detector.html)
+
+### 🔄 Updates
+
+- Change default values for detectors
+
+### 📘 Jupyter Notebooks
+
+- [YoloOnnxDetectorBenchamrks.ipynb](https://github.com/StabRise/ScaleDP-Tutorials/blob/master/object-detection/4.YoloOnnxDetectorBenchmarks.ipynb) - Benchmarking YOLO model with different parameters configurations on CPU and GPU
+
+
 ## [0.2.5] - 10.11.2025
 
 ### 🚀 Features
@@ -14,6 +29,7 @@
 - Improve displaying labels in [ImageDrawBoxes](https://scaledp.stabrise.com/en/latest/image/image_draw_boxes.html)
 
 ### 🧰 Maintenance
+
 - Updated versions of dependencies (Pandas, Numpy, OpenCV)
 
 ### 🐛 Bug Fixes

cliff.toml

@@ -61,6 +61,7 @@ protect_breaking_commits = false
 commit_parsers = [
     { message = "^feat", group = "<!-- 0 -->🚀 Features" },
     { message = "^fix", group = "<!-- 1 -->🐛 Bug Fixes" },
+    { message = "^update", group = "<!-- 2 -->🔄 Updates" },
     { message = "^doc", group = "<!-- 3 -->📚 Documentation" },
     { message = "^perf", group = "<!-- 4 -->⚡ Performance" },
     { message = "^maint", group = "<!-- 4 -->🧰 Maintenance" },

docs/source/release_notes.md

@@ -11,6 +11,21 @@ This document outlines the release notes for the ScaledP project. It includes in
 - Added [TextEmbeddings](#TextEmbeddings) transformer, for compute embedding using SentenceTransformers
 
 
+## [0.2.6] - 19.11.2025
+
+### 🚀 Features
+
+- Enable support GPU in [YoloOnnxDetector](#YoloOnnxDetector)
+
+### 🔄 Updates
+
+- Change default values for detectors
+
+### 📘 Jupyter Notebooks
+
+- [YoloOnnxDetectorBenchamrks.ipynb](https://github.com/StabRise/ScaleDP-Tutorials/blob/master/object-detection/4.YoloOnnxDetectorBenchmarks.ipynb) - Benchmarking YOLO model with different parameters configurations on CPU and GPU
+
+
 ## [0.2.5] - 10.11.2025
 
 ### 🚀 Features

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "scaledp"
-version = "0.3.0rc1"
+version = "0.2.6"
 description = "ScaleDP is a library for processing documents and images using Apache Spark and LLMs"
 authors = ["Mykola Melnyk <mykola@stabrise.com>"]
 repository = "https://github.com/StabRise/scaledp"

scaledp/models/detectors/FaceDetector.py

@@ -18,8 +18,8 @@ class FaceDetector(YoloOnnxDetector):
             "batchSize": 2,
             "partitionMap": False,
             "numPartitions": 0,
-            "pageCol": "page",
-            "pathCol": "path",
+            "pageCol": "",
+            "pathCol": "",
             "propagateError": False,
             "task": "detect",
             "onlyRotated": False,

scaledp/models/detectors/SignatureDetector.py

@@ -16,8 +16,8 @@ class SignatureDetector(YoloOnnxDetector):
             "batchSize": 2,
             "partitionMap": False,
             "numPartitions": 0,
-            "pageCol": "page",
-            "pathCol": "path",
+            "pageCol": "",
+            "pathCol": "",
             "propagateError": False,
             "task": "detect",
             "onlyRotated": False,

scaledp/models/detectors/YoloOnnxDetector.py

@@ -48,8 +48,8 @@ class YoloOnnxDetector(BaseDetector, HasDevice, HasBatchSize, HasLabels):
             "batchSize": 2,
             "partitionMap": False,
             "numPartitions": 0,
-            "pageCol": "page",
-            "pathCol": "path",
+            "pageCol": "",
+            "pathCol": "",
             "propagateError": False,
             "task": "detect",
             "onlyRotated": False,
@@ -83,7 +83,11 @@ def get_model(cls, params):
 
         logging.info("Model downloaded")
 
-        detector = YOLO(model_path_final, conf_thres=params["scoreThreshold"])
+        detector = YOLO(
+            model_path_final,
+            conf_thres=params["scoreThreshold"],
+            device=params["device"],
+        )
 
         cls._model[model_path] = detector
         return cls._model[model_path]

scaledp/models/detectors/yolo/yolo.py

@@ -5,12 +5,13 @@
 import numpy as np
 import onnxruntime
 
+from scaledp.enums import Device
 from scaledp.models.detectors.yolo.utils import multiclass_nms, xywh2xyxy
 
 
 class YOLO:
 
-    def __init__(self, path, conf_thres=0.7, iou_thres=0.5) -> None:
+    def __init__(self, path, device=Device.CPU, conf_thres=0.7, iou_thres=0.5) -> None:
         self.conf_threshold = conf_thres
         self.iou_threshold = iou_thres
 
@@ -22,15 +23,24 @@ def __init__(self, path, conf_thres=0.7, iou_thres=0.5) -> None:
         self.pad_y = None
 
         # Initialize model
-        self.initialize_model(path)
+        self.initialize_model(path, device)
 
     def __call__(self, image) -> Any:
         return self.detect_objects(image)
 
-    def initialize_model(self, path):
-        self.session = onnxruntime.InferenceSession(
-            path, providers=onnxruntime.get_available_providers()
+    def initialize_model(self, path, device):
+        provider = (
+            "CUDAExecutionProvider" if device == Device.CUDA else "CPUExecutionProvider"
         )
+
+        if provider in onnxruntime.get_available_providers():
+            providers = [provider]
+        else:
+            logging.warning(
+                f"{provider} is not available. Falling back to CPUExecutionProvider."
+            )
+            providers = ["CPUExecutionProvider"]
+        self.session = onnxruntime.InferenceSession(path, providers=providers)
         # Get model info
         self.get_input_details()
         self.get_output_details()