Hugging Face's logo

Spaces:
Varun95
/
COCO
Running

App Files Community
COCO / worker.js
Varun95's picture
Varun95
Update worker.js
bf47044 verified
Raw History Blame Contribute Delete
29.6 kB
/*
YOLO11 SEGMENTATION WEB WORKER
Responsibilities:
1. Load ONNX Runtime
2. Initialize WebGPU / WASM
3. Load selected YOLO11n-Seg ONNX model
4. Preprocess camera frames
5. Run inference
6. Decode YOLO segmentation output
7. NMS
8. Generate segmentation masks
*/
import * as ort
from "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.22.0/dist/esm/ort.all.min.mjs";
/* =========================================================
GLOBAL STATE
========================================================= */
let session =
null;
let inputSize =
640;
let inputName =
null;
let outputNames =
[];
/* =========================================================
COCO CLASSES
========================================================= */
const CLASS_NAMES = [
"person",
"bicycle",
"car",
"motorcycle",
"airplane",
"bus",
"train",
"truck",
"boat",
"traffic light",
"fire hydrant",
"stop sign",
"parking meter",
"bench",
"bird",
"cat",
"dog",
"horse",
"sheep",
"cow",
"elephant",
"bear",
"zebra",
"giraffe",
"backpack",
"umbrella",
"handbag",
"tie",
"suitcase",
"frisbee",
"skis",
"snowboard",
"sports ball",
"kite",
"baseball bat",
"baseball glove",
"skateboard",
"surfboard",
"tennis racket",
"bottle",
"wine glass",
"cup",
"fork",
"knife",
"spoon",
"bowl",
"banana",
"apple",
"sandwich",
"orange",
"broccoli",
"carrot",
"hot dog",
"pizza",
"donut",
"cake",
"chair",
"couch",
"potted plant",
"bed",
"dining table",
"toilet",
"tv",
"laptop",
"mouse",
"remote",
"keyboard",
"cell phone",
"microwave",
"oven",
"toaster",
"sink",
"refrigerator",
"book",
"clock",
"vase",
"scissors",
"teddy bear",
"hair drier",
"toothbrush"
];
/* =========================================================
INITIALIZE MODEL
========================================================= */
async function initializeModel(
modelBytes
) {
postMessage({
type:
"loading",
progress:
72,
message:
"Starting ONNX Runtime..."
});
/*
Prefer WebGPU.
*/
let providers =
["wasm"];
let backend =
"WASM";
if (
typeof navigator !==
"undefined" &&
navigator.gpu
) {
providers =
["webgpu"];
backend =
"WebGPU";
}
postMessage({
type:
"loading",
progress:
80,
message:
`Initializing ${backend} backend...`
});
try {
session =
await ort.InferenceSession.create(
modelBytes,
{
executionProviders:
providers,
graphOptimizationLevel:
"all"
}
);
}
catch(error) {
/*
If WebGPU fails,
automatically fall back to WASM.
*/
console.warn(
"WebGPU failed:",
error
);
postMessage({
type:
"loading",
progress:
85,
message:
"WebGPU unavailable. Switching to WASM..."
});
session =
await ort.InferenceSession.create(
modelBytes,
{
executionProviders:
["wasm"],
graphOptimizationLevel:
"all"
}
);
backend =
"WASM";
}
/*
Find model input.
*/
inputName =
session.inputNames[0];
outputNames =
session.outputNames;
console.log(
"Input:",
inputName
);
console.log(
"Outputs:",
outputNames
);
console.log(
"Input metadata:",
session.inputMetadata
);
console.log(
"Output metadata:",
session.outputMetadata
);
postMessage({
type:
"loading",
progress:
94,
message:
"YOLO model loaded. Preparing inference..."
});
/*
Small warm-up inference.
This forces the backend to compile
the graph before the camera starts.
*/
try {
const dummy =
new Float32Array(
1 *
3 *
inputSize *
inputSize
);
const tensor =
new ort.Tensor(
"float32",
dummy,
[
1,
3,
inputSize,
inputSize
]
);
await session.run(
{
[inputName]:
tensor
}
);
}
catch(error) {
console.warn(
"Warmup warning:",
error
);
}
postMessage({
type:
"ready",
backend:
backend
});
}
/* =========================================================
LETTERBOX
========================================================= */
function letterbox(
rgba,
width,
height
) {
const scale =
Math.min(
inputSize / width,
inputSize / height
);
const newWidth =
Math.round(
width * scale
);
const newHeight =
Math.round(
height * scale
);
const padX =
(inputSize - newWidth) / 2;
const padY =
(inputSize - newHeight) / 2;
/*
Create RGB tensor.
YOLO expects:
[1, 3, 640, 640]
RGB
normalized 0-1
*/
const tensorData =
new Float32Array(
3 *
inputSize *
inputSize
);
/*
Gray background.
*/
tensorData.fill(
114 / 255
);
/*
Resize source using nearest-neighbor.
This is intentionally simple and
dependency-free.
*/
for (
let y = 0;
y < newHeight;
y++
) {
const srcY =
Math.min(
height - 1,
Math.floor(
y / scale
)
);
for (
let x = 0;
x < newWidth;
x++
) {
const srcX =
Math.min(
width - 1,
Math.floor(
x / scale
)
);
const srcIndex =
(
srcY *
width +
srcX
) * 4;
const dstX =
Math.floor(
x + padX
);
const dstY =
Math.floor(
y + padY
);
const dstIndex =
dstY *
inputSize +
dstX;
tensorData[
dstIndex
] =
rgba[
srcIndex
] / 255;
tensorData[
inputSize *
inputSize +
dstIndex
] =
rgba[
srcIndex + 1
] / 255;
tensorData[
2 *
inputSize *
inputSize +
dstIndex
] =
rgba[
srcIndex + 2
] / 255;
}
}
return {
tensorData,
scale,
padX,
padY,
newWidth,
newHeight
};
}
/* =========================================================
SIGMOID
========================================================= */
function sigmoid(x) {
return (
1 /
(
1 +
Math.exp(-x)
)
);
}
/* =========================================================
IOU
========================================================= */
function iou(a, b) {
const x1 =
Math.max(
a.x1,
b.x1
);
const y1 =
Math.max(
a.y1,
b.y1
);
const x2 =
Math.min(
a.x2,
b.x2
);
const y2 =
Math.min(
a.y2,
b.y2
);
const intersectionWidth =
Math.max(
0,
x2 - x1
);
const intersectionHeight =
Math.max(
0,
y2 - y1
);
const intersection =
intersectionWidth *
intersectionHeight;
const areaA =
Math.max(
0,
a.x2 - a.x1
) *
Math.max(
0,
a.y2 - a.y1
);
const areaB =
Math.max(
0,
b.x2 - b.x1
) *
Math.max(
0,
b.y2 - b.y1
);
return (
intersection /
(
areaA +
areaB -
intersection +
1e-6
)
);
}
/* =========================================================
NMS
========================================================= */
function nms(
detections,
threshold = 0.45
) {
detections.sort(
(
a,
b
) =>
b.confidence -
a.confidence
);
const selected = [];
while (
detections.length
) {
const current =
detections.shift();
selected.push(
current
);
detections =
detections.filter(
candidate => {
/*
Class-aware NMS.
*/
if (
candidate.classId !==
current.classId
) {
return true;
}
return (
iou(
current,
candidate
) <
threshold
);
}
);
}
return selected;
}
/* =========================================================
FIND OUTPUT TENSORS
========================================================= */
function findDetectionOutput(
results
) {
/*
YOLO11 segmentation normally has:
output:
[1, 116, 8400]
prototype:
[1, 32, 160, 160]
*/
for (
const name of outputNames
) {
const tensor =
results[name];
if (!tensor) {
continue;
}
const dims =
tensor.dims;
if (
dims.length === 3 &&
(
dims[1] === 116 ||
dims[2] === 116
)
) {
return tensor;
}
}
return null;
}
/* =========================================================
FIND PROTOTYPE OUTPUT
========================================================= */
function findPrototypeOutput(
results
) {
for (
const name of outputNames
) {
const tensor =
results[name];
if (!tensor) {
continue;
}
const dims =
tensor.dims;
if (
dims.length === 4 &&
dims[1] === 32
) {
return tensor;
}
}
return null;
}
/* =========================================================
DECODE DETECTIONS
========================================================= */
function decodeDetections(
tensor,
scale,
padX,
padY,
originalWidth,
originalHeight,
confidenceThreshold
) {
const dims =
tensor.dims;
const data =
tensor.data;
/*
Expected:
[1, 116, 8400]
*/
let channels;
let candidates;
let channelFirst =
false;
if (
dims[1] <= 200 &&
dims[2] > 1000
) {
channels =
dims[1];
candidates =
dims[2];
channelFirst =
true;
}
else {
candidates =
dims[1];
channels =
dims[2];
}
const numClasses =
channels - 4 - 32;
const detections = [];
for (
let i = 0;
i < candidates;
i++
) {
/*
Helper to access tensor.
Supports:
[1,C,N]
and
[1,N,C]
*/
function value(
channel
) {
if (
channelFirst
) {
return data[
channel *
candidates +
i
];
}
return data[
i *
channels +
channel
];
}
/*
Find highest class score.
*/
let bestClass =
-1;
let bestScore =
0;
for (
let c = 0;
c < numClasses;
c++
) {
const score =
value(
4 + c
);
if (
score >
bestScore
) {
bestScore =
score;
bestClass =
c;
}
}
if (
bestScore <
confidenceThreshold
) {
continue;
}
/*
YOLO box:
cx
cy
width
height
*/
const cx =
value(0);
const cy =
value(1);
const w =
value(2);
const h =
value(3);
/*
Convert from model coordinates
back to original camera coordinates.
*/
let x1 =
(
cx -
w / 2 -
padX
) / scale;
let y1 =
(
cy -
h / 2 -
padY
) / scale;
let x2 =
(
cx +
w / 2 -
padX
) / scale;
let y2 =
(
cy +
h / 2 -
padY
) / scale;
/*
Clamp
*/
x1 =
Math.max(
0,
Math.min(
originalWidth,
x1
)
);
y1 =
Math.max(
0,
Math.min(
originalHeight,
y1
)
);
x2 =
Math.max(
0,
Math.min(
originalWidth,
x2
)
);
y2 =
Math.max(
0,
Math.min(
originalHeight,
y2
)
);
/*
Mask coefficients.
They start after:
4 box values
+
number of classes
*/
const coefficients =
new Float32Array(
32
);
for (
let m = 0;
m < 32;
m++
) {
coefficients[m] =
value(
4 +
numClasses +
m
);
}
detections.push({
x1,
y1,
x2,
y2,
confidence:
bestScore,
classId:
bestClass,
name:
CLASS_NAMES[
bestClass
] ||
`class-${bestClass}`,
coefficients
});
}
return nms(
detections
);
}
/* =========================================================
CREATE FULL-SIZE MASK
========================================================= */
function createMask(
detection,
prototypeTensor,
originalWidth,
originalHeight,
scale,
padX,
padY
) {
/*
Prototype:
[1, 32, 160, 160]
*/
const protoData =
prototypeTensor.data;
const protoWidth =
prototypeTensor.dims[3];
const protoHeight =
prototypeTensor.dims[2];
const protoChannels =
prototypeTensor.dims[1];
/*
Full original-size binary mask.
0 = background
255 = segmentation
*/
const mask =
new Uint8Array(
originalWidth *
originalHeight
);
/*
Calculate the detection bounding box
in prototype coordinates.
*/
const px1 =
(
detection.x1 *
scale +
padX
) /
inputSize *
protoWidth;
const py1 =
(
detection.y1 *
scale +
padY
) /
inputSize *
protoHeight;
const px2 =
(
detection.x2 *
scale +
padX
) /
inputSize *
protoWidth;
const py2 =
(
detection.y2 *
scale +
padY
) /
inputSize *
protoHeight;
const startX =
Math.max(
0,
Math.floor(px1)
);
const endX =
Math.min(
protoWidth - 1,
Math.ceil(px2)
);
const startY =
Math.max(
0,
Math.floor(py1)
);
const endY =
Math.min(
protoHeight - 1,
Math.ceil(py2)
);
/*
Generate mask.
mask = sigmoid(
coefficient dot prototype
)
*/
for (
let py = startY;
py <= endY;
py++
) {
for (
let px = startX;
px <= endX;
px++
) {
let value =
0;
const protoIndex =
py *
protoWidth +
px;
for (
let c = 0;
c < protoChannels;
c++
) {
value +=
detection
.coefficients[c] *
protoData[
c *
protoWidth *
protoHeight +
protoIndex
];
}
const probability =
sigmoid(
value
);
if (
probability <
0.5
) {
continue;
}
/*
Convert prototype coordinate
to original image coordinate.
*/
const modelX =
(
px /
protoWidth
) *
inputSize;
const modelY =
(
py /
protoHeight
) *
inputSize;
const imageX =
Math.floor(
(
modelX -
padX
) /
scale
);
const imageY =
Math.floor(
(
modelY -
padY
) /
scale
);
if (
imageX >= 0 &&
imageX < originalWidth &&
imageY >= 0 &&
imageY < originalHeight
) {
mask[
imageY *
originalWidth +
imageX
] =
255;
}
}
}
/*
Expand the low-resolution mask
slightly using nearest-neighbor sampling.
This makes the 160x160 prototype
look less sparse on the original frame.
*/
const fullMask =
new Uint8Array(
originalWidth *
originalHeight
);
for (
let y = 0;
y < originalHeight;
y++
) {
const modelY =
Math.floor(
y *
scale
);
const protoY =
Math.floor(
(
modelY +
padY
) /
inputSize *
protoHeight
);
if (
protoY < 0 ||
protoY >= protoHeight
) {
continue;
}
for (
let x = 0;
x < originalWidth;
x++
) {
const modelX =
Math.floor(
x *
scale
);
const protoX =
Math.floor(
(
modelX +
padX
) /
inputSize *
protoWidth
);
if (
protoX < 0 ||
protoX >= protoWidth
) {
continue;
}
/*
Recalculate mask probability
for the current prototype point.
*/
const protoIndex =
protoY *
protoWidth +
protoX;
let value =
0;
for (
let c = 0;
c < protoChannels;
c++
) {
value +=
detection
.coefficients[c] *
protoData[
c *
protoWidth *
protoHeight +
protoIndex
];
}
if (
sigmoid(value) >=
0.5
) {
/*
Restrict to detection box.
*/
if (
x >= detection.x1 &&
x <= detection.x2 &&
y >= detection.y1 &&
y <= detection.y2
) {
fullMask[
y *
originalWidth +
x
] =
255;
}
}
}
}
return fullMask;
}
/* =========================================================
INFERENCE
========================================================= */
async function runInference(
rgbaBuffer,
width,
height,
confidenceThreshold
) {
if (!session) {
throw new Error(
"YOLO model is not initialized."
);
}
/*
Reconstruct RGBA.
*/
const rgba =
new Uint8ClampedArray(
rgbaBuffer
);
/*
Letterbox.
*/
const prepared =
letterbox(
rgba,
width,
height
);
/*
Input tensor.
*/
const inputTensor =
new ort.Tensor(
"float32",
prepared.tensorData,
[
1,
3,
inputSize,
inputSize
]
);
/*
Run YOLO.
*/
const results =
await session.run(
{
[inputName]:
inputTensor
}
);
/*
Find segmentation outputs.
*/
const detectionTensor =
findDetectionOutput(
results
);
const prototypeTensor =
findPrototypeOutput(
results
);
if (!detectionTensor) {
throw new Error(
"Could not find YOLO detection output. " +
"Check the ONNX model output shapes."
);
}
if (!prototypeTensor) {
throw new Error(
"Could not find YOLO segmentation prototype output."
);
}
/*
Decode boxes/classes.
*/
const detections =
decodeDetections(
detectionTensor,
prepared.scale,
prepared.padX,
prepared.padY,
width,
height,
confidenceThreshold
);
/*
Generate masks.
Limit detections to avoid excessive
CPU work on mobile browsers.
*/
const maxMasks =
Math.min(
detections.length,
10
);
for (
let i = 0;
i < maxMasks;
i++
) {
const detection =
detections[i];
try {
const mask =
createMask(
detection,
prototypeTensor,
width,
height,
prepared.scale,
prepared.padX,
prepared.padY
);
detection.mask =
mask.buffer;
detection.maskWidth =
width;
detection.maskHeight =
height;
detection.coefficients =
undefined;
}
catch(error) {
console.warn(
"Mask generation failed:",
error
);
}
}
/*
Remove coefficient arrays
from detections.
*/
for (
const detection
of detections
) {
delete detection.coefficients;
}
/*
Send original camera frame back
so main thread can draw over it.
*/
const frameCopy =
new Uint8ClampedArray(
rgba
);
/*
Transfer buffers to main thread.
*/
const transferList = [
frameCopy.buffer
];
for (
const detection
of detections
) {
if (
detection.mask
) {
transferList.push(
detection.mask
);
}
}
postMessage(
{
type:
"result",
width:
width,
height:
height,
frame:
frameCopy.buffer,
detections:
detections
},
transferList
);
}
/* =========================================================
MESSAGE HANDLER
========================================================= */
self.onmessage =
async function(event) {
const data =
event.data;
/*
MODEL INITIALIZATION
*/
if (
data.type ===
"init"
) {
try {
inputSize =
data.inputSize ||
640;
await initializeModel(
data.modelBytes
);
}
catch(error) {
console.error(
error
);
postMessage({
type:
"error",
message:
"Model initialization failed: " +
error.message
});
}
return;
}
/*
CAMERA FRAME
*/
if (
data.type ===
"frame"
) {
try {
await runInference(
data.rgba,
data.width,
data.height,
data.confidence ||
0.25
);
}
catch(error) {
console.error(
"Inference error:",
error
);
postMessage({
type:
"error",
message:
"Inference error: " +
error.message
});
}
return;
}
};