Removed Activation from ImageSegmenterOptions

This commit is contained in:
kinaryml 2023-04-13 21:09:01 -07:00
parent a03fa448dc
commit a036bf70cc
2 changed files with 2 additions and 12 deletions

View File

@ -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:

View File

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