Merge pull request #80 from StabRise/79-add-support-labels-for-yolo-onnxdetector

79 add support labels for yolo onnxdetector

Author: mykolamelnykml

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
 

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)).

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(),
                 )
 

scaledp/models/detectors/FaceDetector.py

@@ -24,5 +24,6 @@ class FaceDetector(YoloOnnxDetector):
             "task": "detect",
             "onlyRotated": False,
             "model": "StabRise/face_detection",
+            "labels": ["face"],
         },
     )

scaledp/models/detectors/SignatureDetector.py

@@ -22,5 +22,6 @@ class SignatureDetector(YoloOnnxDetector):
             "task": "detect",
             "onlyRotated": False,
             "model": "StabRise/signature_detection",
+            "labels": ["signature"],
         },
     )

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),
             )

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)