chore: Add Image Segmentation binary color thresholding parameter

PiperOrigin-RevId: 726182257
googleapis · Feb 12, 2025 · 0d73929 · 0d73929
1 parent 0aedb1e
commit 0d73929
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 deletions.
diff --git a/tests/unit/aiplatform/test_vision_models.py b/tests/unit/aiplatform/test_vision_models.py
@@ -1280,8 +1280,27 @@ def test_get_image_segmentation_results(self):
             target=prediction_service_client.PredictionServiceClient,
             attribute="predict",
             return_value=gca_prediction_response,
-        ):
-            segmentation_response = model.segment_image(base_image=image)
+        ) as mock_predict:
+            binary_color_threshold = 48
+            segmentation_response = model.segment_image(
+                base_image=image,
+                confidence_threshold=0.1,
+                binary_color_threshold=binary_color_threshold,
+            )
+            mock_predict.assert_called_once_with(
+                endpoint="projects/123456789/locations/us-central1/publishers/google/models/image-segmentation-001",
+                instances=[
+                    {
+                        "base64EncodedImage": image._image_bytes,
+                        "parameters": {
+                            "mode": "foreground",
+                            "confidenceThreshold": 0.1,
+                            "binaryColorThreshold": 48,
+                        },
+                    }
+                ],
+                retry=base._DEFAULT_RETRY,
+            )
             assert len(segmentation_response) == 1
 
 

diff --git a/vertexai/vision_models/_vision_models.py b/vertexai/vision_models/_vision_models.py
@@ -1938,6 +1938,7 @@ def segment_image(
         max_predictions: Optional[int] = None,
         confidence_threshold: Optional[float] = 0.1,
         mask_dilation: Optional[float] = None,
+        binary_color_threshold: Optional[float] = None,
     ) -> ImageSegmentationResponse:
         """Segments an image.
 
@@ -1967,6 +1968,10 @@ def segment_image(
             mask_dilation: A value to dilate the masks by. The value must be in the
                 range of 0.0 (no dilation) and 1.0 (the whole image will be masked).
                 The default is 0.0.
+            binary_color_threshold: The threshold to convert the grayscale soft
+                mask to a binary color black and white mask. The value must be
+                in the range of 0 and 255, or -1 to disable the thresholding.
+                The default is 96.
 
         Returns:
             An `ImageSegmentationResponse` object with the generated masks,
@@ -2000,6 +2005,8 @@ def segment_image(
             parameters["confidenceThreshold"] = confidence_threshold
         if mask_dilation:
             parameters["maskDilation"] = mask_dilation
+        if binary_color_threshold:
+            parameters["binaryColorThreshold"] = binary_color_threshold
 
         response = self._endpoint.predict(
             instances=[instance],