Merge pull request #713 from roboflow/size-measurement-block

Added the Size Measurement Block

inference/core/workflows/core_steps/classical_cv/size_measurement/v1.py

@@ -0,0 +1,155 @@
+from typing import List, Literal, Optional, Tuple, Type, Union
+
+import cv2 as cv
+import numpy as np
+import supervision as sv
+from pydantic import ConfigDict, Field, validator
+
+from inference.core.workflows.execution_engine.entities.base import (
+    Batch,
+    OutputDefinition,
+)
+from inference.core.workflows.execution_engine.entities.types import (
+    INSTANCE_SEGMENTATION_PREDICTION_KIND,
+    LIST_OF_VALUES_KIND,
+    OBJECT_DETECTION_PREDICTION_KIND,
+    STRING_KIND,
+    StepOutputSelector,
+    WorkflowParameterSelector,
+)
+from inference.core.workflows.prototypes.block import (
+    BlockResult,
+    WorkflowBlock,
+    WorkflowBlockManifest,
+)
+
+OUTPUT_KEY: str = "dimensions"
+SHORT_DESCRIPTION = (
+    "Measure the dimensions of objects in relation to a reference object."
+)
+LONG_DESCRIPTION = """
+The `SizeMeasurementBlock` is a transformer block designed to measure the dimensions of objects
+in relation to a reference object. The reference object is detected using one model,
+and the object to be measured is detected using another model. The block outputs the dimensions of the
+objects to be measured in terms of the reference object.
+Note: if reference_predictions provides multiple boxes, the most confident one will be selected.
+In order to achieve different behavior you can use Detection Transformation block with custom filter
+and also continue_if block if no reference detection meets expectations.
+"""
+
+
+class SizeMeasurementManifest(WorkflowBlockManifest):
+    model_config = ConfigDict(
+        json_schema_extra={
+            "name": "Size Measurement",
+            "version": "v1",
+            "short_description": SHORT_DESCRIPTION,
+            "long_description": LONG_DESCRIPTION,
+            "license": "Apache-2.0",
+            "block_type": "transformation",
+        }
+    )
+    type: Literal[f"roboflow_core/size_measurement@v1"]
+    reference_predictions: StepOutputSelector(
+        kind=[
+            INSTANCE_SEGMENTATION_PREDICTION_KIND,
+            OBJECT_DETECTION_PREDICTION_KIND,
+        ]
+    ) = Field(
+        description="Predictions from the reference object model",
+        examples=["$segmentation.reference_predictions"],
+    )
+    object_predictions: StepOutputSelector(
+        kind=[
+            INSTANCE_SEGMENTATION_PREDICTION_KIND,
+            OBJECT_DETECTION_PREDICTION_KIND,
+        ]
+    ) = Field(
+        description="Predictions from the model that detects the object to measure",
+        examples=["$segmentation.object_predictions"],
+    )
+    reference_dimensions: Union[
+        str,
+        Tuple[float, float],
+        List[float],
+        WorkflowParameterSelector(
+            kind=[STRING_KIND, LIST_OF_VALUES_KIND],
+        ),
+    ] = Field(  # type: ignore
+        description="Dimensions of the reference object (width, height) in desired units (e.g., inches) as a string in the format 'width,height' or as a tuple (width, height)",
+        examples=["5.0,5.0", (5.0, 5.0), "$inputs.reference_dimensions"],
+    )
+
+    @classmethod
+    def describe_outputs(cls) -> List[OutputDefinition]:
+        return [
+            OutputDefinition(name=OUTPUT_KEY, kind=[LIST_OF_VALUES_KIND]),
+        ]
+
+    @classmethod
+    def get_execution_engine_compatibility(cls) -> Optional[str]:
+        return ">=1.0.0,<2.0.0"
+
+
+def get_detection_dimensions(
+    detection: sv.Detections, index: int
+) -> Tuple[float, float]:
+    if detection.mask is not None:
+        mask = detection.mask[index].astype(np.uint8)
+        *_, w, h = cv.boundingRect(mask)
+    else:
+        bbox = detection.xyxy[index]
+        w, h = bbox[2] - bbox[0], bbox[3] - bbox[1]
+    return w, h
+
+
+class SizeMeasurementBlockV1(WorkflowBlock):
+    @classmethod
+    def get_manifest(cls) -> Type[WorkflowBlockManifest]:
+        return SizeMeasurementManifest
+
+    def run(
+        self,
+        reference_predictions: sv.Detections,
+        object_predictions: sv.Detections,
+        reference_dimensions: Union[str, Tuple[float, float]],
+    ) -> BlockResult:
+        if isinstance(reference_dimensions, str):
+            try:
+                reference_dimensions = [
+                    float(d) for d in reference_dimensions.split(",")
+                ]
+            except ValueError:
+                raise ValueError(
+                    "reference_dimensions must be a string in the format 'width,height'"
+                )
+        if (
+            not isinstance(reference_dimensions, (tuple, list))
+            or len(reference_dimensions) != 2
+        ):
+            raise ValueError(
+                "reference_dimensions must be a string in the format 'width,height' or a tuple (width, height)"
+            )
+
+        ref_width_actual, ref_height_actual = reference_dimensions
+
+        ref_index = np.argmax(reference_predictions.confidence)
+        ref_width_pixels, ref_height_pixels = get_detection_dimensions(
+            reference_predictions, ref_index
+        )
+
+        if ref_width_pixels == 0 or ref_height_pixels == 0:
+            return {OUTPUT_KEY: None}
+
+        dimensions = []
+        for i in range(len(object_predictions)):
+            obj_width_pixels, obj_height_pixels = get_detection_dimensions(
+                object_predictions, i
+            )
+            obj_width_actual = (obj_width_pixels / ref_width_pixels) * ref_width_actual
+            obj_height_actual = (
+                obj_height_pixels / ref_height_pixels
+            ) * ref_height_actual
+            dimensions.append({"width": obj_width_actual, "height": obj_height_actual})
+
+        return {OUTPUT_KEY: dimensions}

inference/core/workflows/core_steps/loader.py

@@ -39,6 +39,9 @@
 from inference.core.workflows.core_steps.classical_cv.sift_comparison.v2 import (
     SIFTComparisonBlockV2,
 )
+from inference.core.workflows.core_steps.classical_cv.size_measurement.v1 import (
+    SizeMeasurementBlockV1,
+)
 from inference.core.workflows.core_steps.classical_cv.template_matching.v1 import (
     TemplateMatchingBlockV1,
 )
@@ -316,6 +319,7 @@ def load_blocks() -> List[Type[WorkflowBlock]]:
         ContinueIfBlockV1,
         PerspectiveCorrectionBlockV1,
         DynamicZonesBlockV1,
+        SizeMeasurementBlockV1,
         DetectionsClassesReplacementBlockV1,
         ExpressionBlockV1,
         PropertyDefinitionBlockV1,

inference/models/__init__.py

@@ -87,4 +87,8 @@
 )
 from inference.models.yolov9 import YOLOv9ObjectDetection
 from inference.models.yolov10 import YOLOv10ObjectDetection
-from inference.models.yolov11 import YOLOv11ObjectDetection, YOLOv11InstanceSegmentation, YOLOv11KeypointsDetection
+from inference.models.yolov11 import (
+    YOLOv11InstanceSegmentation,
+    YOLOv11KeypointsDetection,
+    YOLOv11ObjectDetection,
+)

inference/models/yolov11/__init__.py

@@ -1,3 +1,7 @@
+from inference.models.yolov11.yolov11_instance_segmentation import (
+    YOLOv11InstanceSegmentation,
+)
+from inference.models.yolov11.yolov11_keypoints_detection import (
+    YOLOv11KeypointsDetection,
+)
 from inference.models.yolov11.yolov11_object_detection import YOLOv11ObjectDetection
-from inference.models.yolov11.yolov11_instance_segmentation import YOLOv11InstanceSegmentation
-from inference.models.yolov11.yolov11_keypoints_detection import YOLOv11KeypointsDetection

tests/workflows/unit_tests/core_steps/classical_cv/test_size_measurement.py

@@ -0,0 +1,83 @@
+import numpy as np
+import pytest
+import supervision as sv
+from inference.core.workflows.core_steps.classical_cv.size_measurement.v1 import (
+    SizeMeasurementBlockV1,
+    OUTPUT_KEY,
+)
+from inference.core.workflows.execution_engine.entities.base import Batch
+
+def test_size_measurement_block():
+    # given
+    reference_predictions = sv.Detections(
+        xyxy=np.array([[10, 10, 50, 50]]),
+        confidence=np.array([0.9]),
+        mask=None,
+        class_id=np.array([0])
+    )
+    object_predictions = sv.Detections(
+        xyxy=np.array([[20, 20, 60, 60], [30, 30, 70, 70]]),
+        confidence=np.array([0.8, 0.85]),
+        mask=None,
+        class_id=np.array([0, 1])
+    )
+    reference_dimensions = "5.0,5.0"
+
+    # when
+    block = SizeMeasurementBlockV1()
+    result = block.run(reference_predictions, object_predictions, reference_dimensions)
+
+    # then
+    expected_dimensions = [
+        {"width": 5.0, "height": 5.0},
+        {"width": 5.0, "height": 5.0}
+    ]
+    assert result == {OUTPUT_KEY: expected_dimensions}, f"Expected {expected_dimensions}, but got {result}"
+
+def test_size_measurement_block_with_mask():
+    # given
+    reference_predictions = sv.Detections(
+        xyxy=np.array([[10, 10, 50, 50]]),
+        confidence=np.array([0.9]),
+        mask=np.array([np.ones((100, 100), dtype=np.uint8)]),
+        class_id=np.array([0])
+    )
+    object_predictions = sv.Detections(
+        xyxy=np.array([[20, 20, 60, 60], [30, 30, 70, 70]]),
+        confidence=np.array([0.8, 0.85]),
+        mask=np.array([np.ones((100, 100), dtype=np.uint8), np.ones((100, 100), dtype=np.uint8)]),
+        class_id=np.array([0, 1])
+    )
+    reference_dimensions = "5.0,5.0"
+
+    # when
+    block = SizeMeasurementBlockV1()
+    result = block.run(reference_predictions, object_predictions, reference_dimensions)
+
+    # then
+    expected_dimensions = [
+        {"width": 5.0, "height": 5.0},
+        {"width": 5.0, "height": 5.0}
+    ]
+    assert result == {OUTPUT_KEY: expected_dimensions}, f"Expected {expected_dimensions}, but got {result}"
+
+def test_size_measurement_block_with_invalid_reference_dimensions():
+    # given
+    reference_predictions = sv.Detections(
+        xyxy=np.array([[10, 10, 50, 50]]),
+        confidence=np.array([0.9]),
+        mask=None,
+        class_id=np.array([0])
+    )
+    object_predictions = sv.Detections(
+        xyxy=np.array([[20, 20, 60, 60], [30, 30, 70, 70]]),
+        confidence=np.array([0.8, 0.85]),
+        mask=None,
+        class_id=np.array([0, 1])
+    )
+    reference_dimensions = "invalid"
+
+    # when
+    block = SizeMeasurementBlockV1()
+    with pytest.raises(expected_exception=ValueError, match="reference_dimensions must be a string in the format 'width,height'"):
+        block.run(reference_predictions, object_predictions, reference_dimensions)