From e070726e697f0fec998c4f334b92a8654405d210 Mon Sep 17 00:00:00 2001 From: mykola Date: Thu, 6 Nov 2025 12:21:22 +0300 Subject: [PATCH 1/4] feat: Added labels param to the YoloOnnxDetector --- scaledp/models/detectors/FaceDetector.py | 1 + scaledp/models/detectors/SignatureDetector.py | 1 + scaledp/models/detectors/YoloOnnxDetector.py | 19 +++++++++---- scaledp/params.py | 28 +++++++++++++++++++ 4 files changed, 44 insertions(+), 5 deletions(-) diff --git a/scaledp/models/detectors/FaceDetector.py b/scaledp/models/detectors/FaceDetector.py index fe4dd1b..85909a0 100644 --- a/scaledp/models/detectors/FaceDetector.py +++ b/scaledp/models/detectors/FaceDetector.py @@ -24,5 +24,6 @@ class FaceDetector(YoloOnnxDetector): "task": "detect", "onlyRotated": False, "model": "StabRise/face_detection", + "labels": ["face"], }, ) diff --git a/scaledp/models/detectors/SignatureDetector.py b/scaledp/models/detectors/SignatureDetector.py index 5d4c8eb..4e5bf98 100644 --- a/scaledp/models/detectors/SignatureDetector.py +++ b/scaledp/models/detectors/SignatureDetector.py @@ -22,5 +22,6 @@ class SignatureDetector(YoloOnnxDetector): "task": "detect", "onlyRotated": False, "model": "StabRise/signature_detection", + "labels": ["signature"], }, ) diff --git a/scaledp/models/detectors/YoloOnnxDetector.py b/scaledp/models/detectors/YoloOnnxDetector.py index df6aff6..44d94b6 100644 --- a/scaledp/models/detectors/YoloOnnxDetector.py +++ b/scaledp/models/detectors/YoloOnnxDetector.py @@ -12,12 +12,12 @@ from scaledp.enums import Device from scaledp.models.detectors.BaseDetector import BaseDetector from scaledp.models.detectors.yolo.yolo import YOLO -from scaledp.params import HasBatchSize, HasDevice +from scaledp.params import HasBatchSize, HasDevice, HasLabels from scaledp.schemas.Box import Box from scaledp.schemas.DetectorOutput import DetectorOutput -class YoloOnnxDetector(BaseDetector, HasDevice, HasBatchSize): +class YoloOnnxDetector(BaseDetector, HasDevice, HasBatchSize, HasLabels): """YOLO ONNX object detector.""" _model: ClassVar = {} @@ -54,6 +54,7 @@ class YoloOnnxDetector(BaseDetector, HasDevice, HasBatchSize): "task": "detect", "onlyRotated": False, "padding": 0, # default padding percent + "labels": [], # default empty labels }, ) @@ -82,7 +83,7 @@ def get_model(cls, params): logging.info("Model downloaded") - detector = YOLO(model_path_final, params["scoreThreshold"]) + detector = YOLO(model_path_final, conf_thres=params["scoreThreshold"]) cls._model[model_path] = detector return cls._model[model_path] @@ -102,7 +103,8 @@ def call_detector(cls, images, params): # Expand boxes by padding percent if provided pad_percent = int(params.get("padding", 0)) if params is not None else 0 h_img, w_img = image_np.shape[:2] - for box in raw_boxes: + labels = params.get("labels", []) + for i, box in enumerate(raw_boxes): # Assume box format is [x1, y1, x2, y2] if pad_percent and len(box) >= 4: x1, y1, x2, y2 = ( @@ -122,7 +124,14 @@ def call_detector(cls, images, params): expanded_box = [x1_new, y1_new, x2_new, y2_new] else: expanded_box = box - boxes.append(Box.from_bbox(expanded_box)) + # Map class_id to label and get score + label = ( + labels[class_ids[i]] + if labels and class_ids[i] < len(labels) + else str(class_ids[i]) + ) + score = scores[i] if scores is not None and i < len(scores) else 0.0 + boxes.append(Box.from_bbox(expanded_box, label=label, score=score)) results_final.append( DetectorOutput(path=image_path, type="yolo", bboxes=boxes), ) diff --git a/scaledp/params.py b/scaledp/params.py index ea1bca8..7019c4d 100644 --- a/scaledp/params.py +++ b/scaledp/params.py @@ -827,3 +827,31 @@ def getPropagateError(self) -> bool: def setPropagateError(self, value: bool) -> Any: """Sets the value of :py:attr:`propagateError`.""" return self._set(propagateError=value) + + +class HasLabels(Params): + """ + Mixin for param labels: list of class labels. + """ + + labels: "Param[list[str]]" = Param( + Params._dummy(), + "labels", + "List of the labels.", + typeConverter=TypeConverters.toListString, + ) + + def __init__(self) -> None: + super(HasLabels, self).__init__() + + def getLabels(self) -> list[str]: + """ + Gets the value of labels or its default value. + """ + return self.getOrDefault(self.labels) + + def setLabels(self, value: list[str]) -> Any: + """ + Sets the value of :py:attr:`labels`. + """ + return self._set(labels=value) From def99a8202d0580b7178891fd81a2772277eeb0d Mon Sep 17 00:00:00 2001 From: mykola Date: Thu, 6 Nov 2025 12:25:14 +0300 Subject: [PATCH 2/4] feat: Improve displaying labels in ImageDrawBoxes.py --- scaledp/image/ImageDrawBoxes.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/scaledp/image/ImageDrawBoxes.py b/scaledp/image/ImageDrawBoxes.py index f7ec526..aaa3961 100644 --- a/scaledp/image/ImageDrawBoxes.py +++ b/scaledp/image/ImageDrawBoxes.py @@ -179,13 +179,33 @@ def draw_boxes(self, data, fill, img1): self.draw_box(box, color, fill, img1) text = self.getDisplayText(box) if text: + tbox = list( + img1.textbbox( + ( + box.x, + box.y - self.getTextSize() * 1.2 - self.getPadding(), + ), + text, + font_size=self.getTextSize(), + ), + ) + tbox[3] = tbox[3] + self.getTextSize() / 4 + tbox[2] = tbox[2] + self.getTextSize() / 4 + tbox[0] = box.x - self.getPadding() + tbox[1] = box.y - self.getTextSize() * 1.2 - self.getPadding() + img1.rounded_rectangle( + tbox, + outline=color, + radius=2, + fill=color, + ) img1.text( ( box.x, box.y - self.getTextSize() * 1.2 - self.getPadding(), ), text, - fill=color, + fill="white", font_size=self.getTextSize(), ) From b655a7a744213003fb1b538d47403c4cce129b65 Mon Sep 17 00:00:00 2001 From: mykola Date: Thu, 6 Nov 2025 12:30:09 +0300 Subject: [PATCH 3/4] Updated CHANGELOG.md --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6e84687..142bd9e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,8 @@ ### 🚀 Features - Added param 'returnEmpty' to [ImageCropBoxes](https://scaledp.stabrise.com/en/latest/image/image_crop_boxes.html) for avoid to have exceptions if no boxes are found +- Added labels param to the [YoloOnnxDetector](https://scaledp.stabrise.com/en/latest/models/detectors/yolo_onnx_detector.html) +- Improve displaying labels in [ImageDrawBoxes](https://scaledp.stabrise.com/en/latest/image/image_draw_boxes.html) ### 🐛 Bug Fixes From 9287d29642824a7101fb039e6504e6daa7c34585 Mon Sep 17 00:00:00 2001 From: mykola Date: Thu, 6 Nov 2025 20:28:53 +0300 Subject: [PATCH 4/4] Updated docs --- docs/source/models/detectors/yolo_onnx_detector.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/source/models/detectors/yolo_onnx_detector.md b/docs/source/models/detectors/yolo_onnx_detector.md index 751e39e..254f2ce 100644 --- a/docs/source/models/detectors/yolo_onnx_detector.md +++ b/docs/source/models/detectors/yolo_onnx_detector.md @@ -47,11 +47,12 @@ detected_df = detector.transform(input_df) | onlyRotated | bool | Return only rotated boxes | False | | model | str | Model identifier | (required) | | padding | int | Padding percent to expand detected boxes | 0 | +| labels | list | List of class labels for detected objects | [] | ## Notes - The detector loads YOLO ONNX models from Hugging Face Hub or local path. - Supports batch and distributed processing with Spark. - Padding expands detected bounding boxes by a percentage. +- The `labels` parameter allows you to specify custom class labels for the detected objects. If provided, detection results will use these labels instead of class indices. - Used as a base for specialized detectors (e.g., [**Face Detector**] - (#FaceDetector), [**Signature Detector**](#SignatureDetector)). - + (#FaceDetector), [**Signature Detector**](#SignatureDetector)). \ No newline at end of file