ai-tube / src /lib /on-device-ai /segmentFrameOnClick.ts
jbilcke-hf's picture
jbilcke-hf HF staff
some more craziness
fddab62
raw
history blame
No virus
1.88 kB
import { FilesetResolver, InteractiveSegmenter, InteractiveSegmenterResult } from "@mediapipe/tasks-vision"
export type InteractiveVideoSegmenter = (videoFrame: TexImageSource, x: number, y: number) => Promise<InteractiveSegmenterResult>
const getInteractiveSegmenter = async (): Promise<InteractiveVideoSegmenter> => {
const vision = await FilesetResolver.forVisionTasks(
"https://cdn.jsdelivr.net/npm/@mediapipe/tasks-vision@latest/wasm"
);
const interactiveSegmenter = await InteractiveSegmenter.createFromOptions(vision, {
baseOptions: {
modelAssetPath:
"https://storage.googleapis.com/mediapipe-tasks/interactive_segmenter/ptm_512_hdt_ptm_woid.tflite"
},
outputCategoryMask: true,
outputConfidenceMasks: false,
});
const segmenter: InteractiveVideoSegmenter = (
videoFrame: TexImageSource,
x: number,
y: number
): Promise<InteractiveSegmenterResult> => {
return new Promise((resolve, reject) => {
interactiveSegmenter.segment(
videoFrame,
{
keypoint: { x, y }
},
(results) => {
resolve(results)
})
})
}
return segmenter
}
const globalState: { segmenter?: InteractiveVideoSegmenter } = {};
(async () => {
globalState.segmenter = globalState.segmenter || (await getInteractiveSegmenter())
})();
export async function segmentFrame(frame: TexImageSource, x: number, y: number): Promise<InteractiveSegmenterResult> {
console.log("segmentFrame: loading segmenter..")
globalState.segmenter = globalState.segmenter || (await getInteractiveSegmenter())
console.log("segmentFrame: segmenting..")
return globalState.segmenter(frame, x, y)
}
// to run:
// see doc:
// https://developers.google.com/mediapipe/solutions/vision/image_segmenter/web_js#video
// imageSegmenter.segmentForVideo(video, startTimeMs, callbackForVideo);