Use Uint8ClampedArray for pixel output

PiperOrigin-RevId: 518362677
This commit is contained in:
Sebastian Schmidt 2023-03-21 13:19:44 -07:00 committed by Copybara-Service
parent 7d26daf723
commit a5fc1d4baf
5 changed files with 15 additions and 16 deletions

View File

@ -36,9 +36,8 @@ const COLOR_MAP = [
/** Helper function to draw a confidence mask */ /** Helper function to draw a confidence mask */
export function drawConfidenceMask( export function drawConfidenceMask(
ctx: CanvasRenderingContext2D, image: Float32Array, width: number,
ctx: CanvasRenderingContext2D, image: Float32Array|Uint8Array, height: number): void {
width: number, height: number): void {
const uint8ClampedArray = new Uint8ClampedArray(width * height * 4); const uint8ClampedArray = new Uint8ClampedArray(width * height * 4);
for (let i = 0; i < image.length; i++) { for (let i = 0; i < image.length; i++) {
uint8ClampedArray[4 * i] = 128; uint8ClampedArray[4 * i] = 128;
@ -54,9 +53,9 @@ export function drawConfidenceMask(
* for now. * for now.
*/ */
export function drawCategoryMask( export function drawCategoryMask(
ctx: CanvasRenderingContext2D, image: Float32Array|Uint8Array, ctx: CanvasRenderingContext2D, image: Uint8ClampedArray|Float32Array,
width: number, height: number): void { width: number, height: number): void {
const uint8ClampedArray = new Uint8ClampedArray(width * height * 4); const rgbaArray = new Uint8ClampedArray(width * height * 4);
const isFloatArray = image instanceof Float32Array; const isFloatArray = image instanceof Float32Array;
for (let i = 0; i < image.length; i++) { for (let i = 0; i < image.length; i++) {
const colorIndex = isFloatArray ? Math.round(image[i] * 255) : image[i]; const colorIndex = isFloatArray ? Math.round(image[i] * 255) : image[i];
@ -69,10 +68,10 @@ export function drawCategoryMask(
return; return;
} }
uint8ClampedArray[4 * i] = color[0]; rgbaArray[4 * i] = color[0];
uint8ClampedArray[4 * i + 1] = color[1]; rgbaArray[4 * i + 1] = color[1];
uint8ClampedArray[4 * i + 2] = color[2]; rgbaArray[4 * i + 2] = color[2];
uint8ClampedArray[4 * i + 3] = color[3]; rgbaArray[4 * i + 3] = color[3];
} }
ctx.putImageData(new ImageData(uint8ClampedArray, width, height), 0, 0); ctx.putImageData(new ImageData(rgbaArray, width, height), 0, 0);
} }

View File

@ -17,17 +17,17 @@
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 Uint8Array * The segmentation tasks return the segmentation result as a Uint8ClampedArray
* (when the default mode of `CATEGORY_MASK` is used) or as a Float32Array (for * (when the default mode of `CATEGORY_MASK` is used) or as a Float32Array (for
* output type `CONFIDENCE_MASK`). The `WebGLTexture` output type is reserved * output type `CONFIDENCE_MASK`). The `WebGLTexture` output type is reserved
* for future usage. * for future usage.
*/ */
export type SegmentationMask = Uint8Array|Float32Array|WebGLTexture; export type SegmentationMask = Uint8ClampedArray|Float32Array|WebGLTexture;
/** /**
* A callback that receives the computed masks from the segmentation tasks. The * A callback that receives the computed masks from the segmentation tasks. The
* callback either receives a single element array with a category mask (as a * callback either receives a single element array with a category mask (as a
* `[Uint8Array]`) or multiple confidence masks (as a `Float32Array[]`). * `[Uint8ClampedArray]`) or multiple confidence masks (as a `Float32Array[]`).
* The returned data is only valid for the duration of the callback. If * The returned data is only valid for the duration of the callback. If
* asynchronous processing is needed, all data needs to be copied before the * asynchronous processing is needed, all data needs to be copied before the
* callback returns. * callback returns.

View File

@ -159,7 +159,7 @@ describe('ImageSegmenter', () => {
}); });
it('supports category masks', (done) => { it('supports category masks', (done) => {
const mask = new Uint8Array([1, 2, 3, 4]); const mask = new Uint8ClampedArray([1, 2, 3, 4]);
// Pass the test data to our listener // Pass the test data to our listener
imageSegmenter.fakeWasmModule._waitUntilIdle.and.callFake(() => { imageSegmenter.fakeWasmModule._waitUntilIdle.and.callFake(() => {

View File

@ -154,7 +154,7 @@ describe('InteractiveSegmenter', () => {
}); });
it('supports category masks', (done) => { it('supports category masks', (done) => {
const mask = new Uint8Array([1, 2, 3, 4]); const mask = new Uint8ClampedArray([1, 2, 3, 4]);
// Pass the test data to our listener // Pass the test data to our listener
interactiveSegmenter.fakeWasmModule._waitUntilIdle.and.callFake(() => { interactiveSegmenter.fakeWasmModule._waitUntilIdle.and.callFake(() => {

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: Uint8Array|Uint8ClampedArray|Float32Array; data: Uint8ClampedArray|Float32Array;
width: number; width: number;
height: number; height: number;
} }