Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
import React from "react" | |
import { ImageSegmenterResult, ImageSource } from "@mediapipe/tasks-vision" | |
import { segmentFrame } from "./segmentFrame" | |
export async function getSegmentationCanvas({ | |
frame, | |
timestamp, | |
width, | |
height | |
}: { | |
frame: ImageSource; | |
timestamp: number; | |
width: number; | |
height: number; | |
}): Promise<JSX.Element> { | |
const results: ImageSegmenterResult = await segmentFrame(frame, timestamp); | |
const canvas: HTMLCanvasElement | OffscreenCanvas | undefined = results.categoryMask?.canvas; | |
// If there is a canvas, style it and return | |
if (canvas) { | |
const style = { | |
width: `${width}px`, | |
height: `${height}px`, | |
}; | |
const CanvasComponent = () => ( | |
<canvas | |
ref={(node) => { | |
if (node) { | |
node.width = width; | |
node.height = height; | |
const context = node.getContext('2d'); | |
if (context) { | |
context.drawImage(canvas, 0, 0, width, height); | |
} | |
} | |
}} | |
style={style} | |
/> | |
); | |
return <CanvasComponent />; | |
} else { | |
// Return a blank canvas if no canvas is found in results | |
return ( | |
<canvas | |
width={width} | |
height={height} | |
style={{ | |
width: `${width}px`, | |
height: `${height}px`, | |
backgroundColor: 'transparent', | |
}} | |
/> | |
); | |
} | |
} |