diff --git a/mediapipe/tasks/web/vision/image_segmenter/image_segmenter.ts b/mediapipe/tasks/web/vision/image_segmenter/image_segmenter.ts index 92462815..60b96534 100644 --- a/mediapipe/tasks/web/vision/image_segmenter/image_segmenter.ts +++ b/mediapipe/tasks/web/vision/image_segmenter/image_segmenter.ts @@ -60,6 +60,7 @@ export type ImageSegmenterCallback = (result: ImageSegmenterResult) => void; export class ImageSegmenter extends VisionTaskRunner { private result: ImageSegmenterResult = {}; private labels: string[] = []; + private userCallback: ImageSegmenterCallback = () => {}; private outputCategoryMask = DEFAULT_OUTPUT_CATEGORY_MASK; private outputConfidenceMasks = DEFAULT_OUTPUT_CONFIDENCE_MASKS; private readonly options: ImageSegmenterGraphOptionsProto; @@ -232,14 +233,13 @@ export class ImageSegmenter extends VisionTaskRunner { typeof imageProcessingOptionsOrCallback !== 'function' ? imageProcessingOptionsOrCallback : {}; - const userCallback = - typeof imageProcessingOptionsOrCallback === 'function' ? + this.userCallback = typeof imageProcessingOptionsOrCallback === 'function' ? imageProcessingOptionsOrCallback : callback!; this.reset(); this.processImageData(image, imageProcessingOptions); - userCallback(this.result); + this.userCallback = () => {}; } /** @@ -286,13 +286,13 @@ export class ImageSegmenter extends VisionTaskRunner { const timestamp = typeof timestampOrImageProcessingOptions === 'number' ? timestampOrImageProcessingOptions : timestampOrCallback as number; - const userCallback = typeof timestampOrCallback === 'function' ? + this.userCallback = typeof timestampOrCallback === 'function' ? timestampOrCallback : callback!; this.reset(); this.processVideoData(videoFrame, imageProcessingOptions, timestamp); - userCallback(this.result); + this.userCallback = () => {}; } /** @@ -314,6 +314,18 @@ export class ImageSegmenter extends VisionTaskRunner { this.result = {}; } + /** Invokes the user callback once all data has been received. */ + private maybeInvokeCallback(): void { + if (this.outputConfidenceMasks && !('confidenceMasks' in this.result)) { + return; + } + if (this.outputCategoryMask && !('categoryMask' in this.result)) { + return; + } + + this.userCallback(this.result); + } + /** Updates the MediaPipe graph configuration. */ protected override refreshGraph(): void { const graphConfig = new CalculatorGraphConfig(); @@ -342,10 +354,13 @@ export class ImageSegmenter extends VisionTaskRunner { this.result.confidenceMasks = masks.map(wasmImage => this.convertToMPImage(wasmImage)); this.setLatestOutputTimestamp(timestamp); + this.maybeInvokeCallback(); }); this.graphRunner.attachEmptyPacketListener( CONFIDENCE_MASKS_STREAM, timestamp => { + this.result.confidenceMasks = undefined; this.setLatestOutputTimestamp(timestamp); + this.maybeInvokeCallback(); }); } @@ -357,10 +372,13 @@ export class ImageSegmenter extends VisionTaskRunner { CATEGORY_MASK_STREAM, (mask, timestamp) => { this.result.categoryMask = this.convertToMPImage(mask); this.setLatestOutputTimestamp(timestamp); + this.maybeInvokeCallback(); }); this.graphRunner.attachEmptyPacketListener( CATEGORY_MASK_STREAM, timestamp => { + this.result.categoryMask = undefined; this.setLatestOutputTimestamp(timestamp); + this.maybeInvokeCallback(); }); } diff --git a/mediapipe/tasks/web/vision/image_segmenter/image_segmenter_test.ts b/mediapipe/tasks/web/vision/image_segmenter/image_segmenter_test.ts index 7f7f1906..c1ccd799 100644 --- a/mediapipe/tasks/web/vision/image_segmenter/image_segmenter_test.ts +++ b/mediapipe/tasks/web/vision/image_segmenter/image_segmenter_test.ts @@ -262,4 +262,34 @@ describe('ImageSegmenter', () => { }); }); }); + + it('invokes listener once masks are avaiblae', async () => { + const categoryMask = new Uint8ClampedArray([1]); + const confidenceMask = new Float32Array([0.0]); + let listenerCalled = false; + + await imageSegmenter.setOptions( + {outputCategoryMask: true, outputConfidenceMasks: true}); + + // Pass the test data to our listener + imageSegmenter.fakeWasmModule._waitUntilIdle.and.callFake(() => { + expect(listenerCalled).toBeFalse(); + imageSegmenter.categoryMaskListener! + ({data: categoryMask, width: 1, height: 1}, 1337); + expect(listenerCalled).toBeFalse(); + imageSegmenter.confidenceMasksListener!( + [ + {data: confidenceMask, width: 1, height: 1}, + ], + 1337); + expect(listenerCalled).toBeTrue(); + }); + + return new Promise(resolve => { + imageSegmenter.segment({} as HTMLImageElement, () => { + listenerCalled = true; + resolve(); + }); + }); + }); });