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 */
export function drawConfidenceMask(
ctx: CanvasRenderingContext2D, image: Float32Array|Uint8Array,
width: number, height: number): void {
ctx: CanvasRenderingContext2D, image: Float32Array, width: number,
height: number): void {
const uint8ClampedArray = new Uint8ClampedArray(width * height * 4);
for (let i = 0; i < image.length; i++) {
uint8ClampedArray[4 * i] = 128;
@ -54,9 +53,9 @@ export function drawConfidenceMask(
* for now.
*/
export function drawCategoryMask(
ctx: CanvasRenderingContext2D, image: Float32Array|Uint8Array,
ctx: CanvasRenderingContext2D, image: Uint8ClampedArray|Float32Array,
width: number, height: number): void {
const uint8ClampedArray = new Uint8ClampedArray(width * height * 4);
const rgbaArray = new Uint8ClampedArray(width * height * 4);
const isFloatArray = image instanceof Float32Array;
for (let i = 0; i < image.length; i++) {
const colorIndex = isFloatArray ? Math.round(image[i] * 255) : image[i];
@ -69,10 +68,10 @@ export function drawCategoryMask(
return;
}
uint8ClampedArray[4 * i] = color[0];
uint8ClampedArray[4 * i + 1] = color[1];
uint8ClampedArray[4 * i + 2] = color[2];
uint8ClampedArray[4 * i + 3] = color[3];
rgbaArray[4 * i] = color[0];
rgbaArray[4 * i + 1] = color[1];
rgbaArray[4 * i + 2] = color[2];
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';
/**
* 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
* output type `CONFIDENCE_MASK`). The `WebGLTexture` output type is reserved
* 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
* 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
* asynchronous processing is needed, all data needs to be copied before the
* callback returns.

View File

@ -159,7 +159,7 @@ describe('ImageSegmenter', () => {
});
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
imageSegmenter.fakeWasmModule._waitUntilIdle.and.callFake(() => {

View File

@ -154,7 +154,7 @@ describe('InteractiveSegmenter', () => {
});
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
interactiveSegmenter.fakeWasmModule._waitUntilIdle.and.callFake(() => {

View File

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