Add WebGLTexture output for ImageSegmenter

PiperOrigin-RevId: 518886135
This commit is contained in:
Sebastian Schmidt 2023-03-23 09:45:25 -07:00 committed by Copybara-Service
parent 1c9e6894f3
commit 8e5eadbd4e
3 changed files with 7 additions and 13 deletions

View File

@ -325,14 +325,7 @@ absl::Status TensorsToSegmentationCalculator::Process(
postprocessor_.GetSegmentationResultGpu(input_shape, output_shape, postprocessor_.GetSegmentationResultGpu(input_shape, output_shape,
input_tensor); input_tensor);
for (int i = 0; i < segmented_masks.size(); ++i) { for (int i = 0; i < segmented_masks.size(); ++i) {
// Real output on GPU. kSegmentationOut(cc)[i].Send(std::move(segmented_masks[i]));
// 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));
} }
return absl::OkStatus(); return absl::OkStatus();
} }

View File

@ -17,10 +17,11 @@
import {NormalizedKeypoint} from '../../../../tasks/web/components/containers/keypoint'; import {NormalizedKeypoint} from '../../../../tasks/web/components/containers/keypoint';
/** /**
* The segmentation tasks return the segmentation result as a Uint8ClampedArray * The segmentation tasks return the segmentation either as a WebGLTexture (when
* (when the default mode of `CATEGORY_MASK` is used) or as a Float32Array (for * the output is on GPU) or as a typed JavaScript arrays for CPU-based
* output type `CONFIDENCE_MASK`). The `WebGLTexture` output type is reserved * category or confidence masks. `Uint8ClampedArray`s are used to represend
* for future usage. * CPU-based category masks and `Float32Array`s are used for CPU-based
* confidence masks.
*/ */
export type SegmentationMask = Uint8ClampedArray|Float32Array|WebGLTexture; export type SegmentationMask = Uint8ClampedArray|Float32Array|WebGLTexture;

View File

@ -10,7 +10,7 @@ type LibConstructor = new (...args: any[]) => GraphRunner;
/** An image returned from a MediaPipe graph. */ /** An image returned from a MediaPipe graph. */
export interface WasmImage { export interface WasmImage {
data: Uint8ClampedArray|Float32Array; data: Uint8ClampedArray|Float32Array|WebGLTexture;
width: number; width: number;
height: number; height: number;
} }