diff --git a/mediapipe/tasks/cc/vision/image_segmenter/calculators/tensors_to_segmentation_calculator.cc b/mediapipe/tasks/cc/vision/image_segmenter/calculators/tensors_to_segmentation_calculator.cc
index b6c1fe6b..0cdc8fe0 100644
--- a/mediapipe/tasks/cc/vision/image_segmenter/calculators/tensors_to_segmentation_calculator.cc
+++ b/mediapipe/tasks/cc/vision/image_segmenter/calculators/tensors_to_segmentation_calculator.cc
@@ -325,14 +325,7 @@ absl::Status TensorsToSegmentationCalculator::Process(
         postprocessor_.GetSegmentationResultGpu(input_shape, output_shape,
                                                 input_tensor);
     for (int i = 0; i < segmented_masks.size(); ++i) {
-      // Real output on GPU.
-      // kSegmentationOut(cc)[i].Send(std::move(segmented_masks[i]));
-
-      // Reformat as CPU for now for testing.
-      // TODO: Switch to real GPU output when GPU output pipeline is
-      //     ready.
-      Image new_image(segmented_masks[i]->GetImageFrameSharedPtr());
-      kSegmentationOut(cc)[i].Send(std::move(new_image));
+      kSegmentationOut(cc)[i].Send(std::move(segmented_masks[i]));
     }
     return absl::OkStatus();
   }
diff --git a/mediapipe/tasks/web/vision/core/types.d.ts b/mediapipe/tasks/web/vision/core/types.d.ts
index f0ac0862..c04366ac 100644
--- a/mediapipe/tasks/web/vision/core/types.d.ts
+++ b/mediapipe/tasks/web/vision/core/types.d.ts
@@ -17,10 +17,11 @@
 import {NormalizedKeypoint} from '../../../../tasks/web/components/containers/keypoint';
 
 /**
- * The segmentation tasks return the segmentation result as a Uint8ClampedArray
- * (when the default mode of `CATEGORY_MASK` is used) or as a Float32Array (for
- * output type `CONFIDENCE_MASK`). The `WebGLTexture` output type is reserved
- * for future usage.
+ * The segmentation tasks return the segmentation either as a WebGLTexture (when
+ * the output is on GPU) or as a typed JavaScript arrays for CPU-based
+ * category or confidence masks. `Uint8ClampedArray`s are used to represend
+ * CPU-based category masks and `Float32Array`s are used for CPU-based
+ * confidence masks.
  */
 export type SegmentationMask = Uint8ClampedArray|Float32Array|WebGLTexture;
 
diff --git a/mediapipe/web/graph_runner/graph_runner_image_lib.ts b/mediapipe/web/graph_runner/graph_runner_image_lib.ts
index d9bb0568..8b491d89 100644
--- a/mediapipe/web/graph_runner/graph_runner_image_lib.ts
+++ b/mediapipe/web/graph_runner/graph_runner_image_lib.ts
@@ -10,7 +10,7 @@ type LibConstructor = new (...args: any[]) => GraphRunner;
 
 /** An image returned from a MediaPipe graph. */
 export interface WasmImage {
-  data: Uint8ClampedArray|Float32Array;
+  data: Uint8ClampedArray|Float32Array|WebGLTexture;
   width: number;
   height: number;
 }