chore: Add Image Segmentation binary color thresholding parameter

PiperOrigin-RevId: 726182257

Author: vertex-sdk-bot

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
 
 

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