diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index aa557281f..de0a263f3 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -15,7 +15,6 @@ import enum import os -from typing import List from unittest import mock from absl.testing import absltest @@ -34,7 +33,6 @@ ImageSegmenterResult = image_segmenter.ImageSegmenterResult _BaseOptions = base_options_module.BaseOptions _Image = image_module.Image _ImageFormat = image_frame.ImageFormat -_Activation = image_segmenter.ImageSegmenterOptions.Activation _ImageSegmenter = image_segmenter.ImageSegmenter _ImageSegmenterOptions = image_segmenter.ImageSegmenterOptions _RUNNING_MODE = vision_task_running_mode.VisionTaskRunningMode @@ -191,7 +189,7 @@ class ImageSegmenterTest(parameterized.TestCase): # Run segmentation on the model in CONFIDENCE_MASK mode. options = _ImageSegmenterOptions( base_options=base_options, output_category_mask=False, - output_confidence_masks=True, activation=_Activation.SOFTMAX + output_confidence_masks=True ) with _ImageSegmenter.create_from_options(options) as segmenter: diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index 102773173..c5204db47 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -78,16 +78,10 @@ class ImageSegmenterOptions: is set to the live stream mode. """ - class Activation(enum.Enum): - NONE = 0 - SIGMOID = 1 - SOFTMAX = 2 - base_options: _BaseOptions running_mode: _RunningMode = _RunningMode.IMAGE output_confidence_masks: bool = True output_category_mask: bool = False - activation: Optional[Activation] = Activation.NONE result_callback: Optional[ Callable[[ImageSegmenterResult, image_module.Image, int], None] ] = None @@ -99,9 +93,7 @@ class ImageSegmenterOptions: base_options_proto.use_stream_mode = ( False if self.running_mode == _RunningMode.IMAGE else True ) - segmenter_options_proto = _SegmenterOptionsProto( - activation=self.activation.value - ) + segmenter_options_proto = _SegmenterOptionsProto() return _ImageSegmenterGraphOptionsProto( base_options=base_options_proto, segmenter_options=segmenter_options_proto,