From 35a04522faecd3e82acedeeaaf747da2615820c4 Mon Sep 17 00:00:00 2001
From: kinaryml <kinar.ravishankar@ymedialabs.com>
Date: Thu, 3 Nov 2022 00:46:00 -0700
Subject: [PATCH] Moved the OutputType and Activation classes to
 ImageSegmenter's inner classes

---
 .../test/vision/image_segmenter_test.py       |  4 ++--
 .../tasks/python/vision/image_segmenter.py    | 23 +++++++++----------
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py
index dde751c1d..5072d3482 100644
--- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py
+++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py
@@ -33,8 +33,8 @@ from mediapipe.tasks.python.vision.core import vision_task_running_mode
 _BaseOptions = base_options_module.BaseOptions
 _Image = image_module.Image
 _ImageFormat = image_frame.ImageFormat
-_OutputType = image_segmenter.OutputType
-_Activation = image_segmenter.Activation
+_OutputType = image_segmenter.ImageSegmenterOptions.OutputType
+_Activation = image_segmenter.ImageSegmenterOptions.Activation
 _ImageSegmenter = image_segmenter.ImageSegmenter
 _ImageSegmenterOptions = image_segmenter.ImageSegmenterOptions
 _RUNNING_MODE = vision_task_running_mode.VisionTaskRunningMode
diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py
index ebfeca1bc..c1b50a5ae 100644
--- a/mediapipe/tasks/python/vision/image_segmenter.py
+++ b/mediapipe/tasks/python/vision/image_segmenter.py
@@ -44,18 +44,6 @@ _TASK_GRAPH_NAME = 'mediapipe.tasks.vision.ImageSegmenterGraph'
 _MICRO_SECONDS_PER_MILLISECOND = 1000
 
 
-class OutputType(enum.Enum):
-  UNSPECIFIED = 0
-  CATEGORY_MASK = 1
-  CONFIDENCE_MASK = 2
-
-
-class Activation(enum.Enum):
-  NONE = 0
-  SIGMOID = 1
-  SOFTMAX = 2
-
-
 @dataclasses.dataclass
 class ImageSegmenterOptions:
   """Options for the image segmenter task.
@@ -74,6 +62,17 @@ class ImageSegmenterOptions:
       data. The result callback should only be specified when the running mode
       is set to the live stream mode.
   """
+
+  class OutputType(enum.Enum):
+    UNSPECIFIED = 0
+    CATEGORY_MASK = 1
+    CONFIDENCE_MASK = 2
+
+  class Activation(enum.Enum):
+    NONE = 0
+    SIGMOID = 1
+    SOFTMAX = 2
+
   base_options: _BaseOptions
   running_mode: _RunningMode = _RunningMode.IMAGE
   output_type: Optional[OutputType] = OutputType.CATEGORY_MASK