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,
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();
}

View File

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

View File

@ -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;
}