ONNX model post process logic fail to crack
question
In this code i segmentation mask i not able load correctly it just giving border of image like :
but i want the segmentation mask like ( detect back side part of image) :
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>YOLOv8 Segmentation</title>
<script src="https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.min.js"></script>
<style>
canvas {
display: block;
border: 1px solid black;
margin-top: 10px;
}
</style>
</head>
<body>
<input id="uploadInput" type="file" />
<canvas></canvas>
<script>
const input = document.getElementById("uploadInput");
input.addEventListener("change", async (event) => {
await detect_and_draw_segmentation(event.target.files[0]);
});
async function detect_and_draw_segmentation(file) {
const img = new Image();
img.src = URL.createObjectURL(file);
await img.decode();
const canvas = document.querySelector("canvas");
canvas.width = img.width;
canvas.height = img.height;
const ctx = canvas.getContext("2d");
// Resize and normalize image to 640x640
const tmpCanvas = document.createElement("canvas");
tmpCanvas.width = 640;
tmpCanvas.height = 640;
const tmpCtx = tmpCanvas.getContext("2d");
tmpCtx.drawImage(img, 0, 0, 640, 640);
const imgData = tmpCtx.getImageData(0, 0, 640, 640);
const data = imgData.data;
// Normalize to CHW (1, 3, 640, 640)
const red = [], green = [], blue = [];
for (let i = 0; i < data.length; i += 4) {
red.push(data[i] / 255);
green.push(data[i + 1] / 255);
blue.push(data[i + 2] / 255);
}
const inputTensor = new ort.Tensor(
"float32",
Float32Array.from([...red, ...green, ...blue]),
[1, 3, 640, 640]
);
// Load ONNX model
const session = await ort.InferenceSession.create("best_on.onnx");
const feeds = { images: inputTensor };
const results = await session.run(feeds);
const output0 = results["output0"].data; // [batch, num_boxes, 4+num_classes+32] (e.g., 8400x117)
const proto = results["output1"].data; // [1, 32, 160, 160]
const num_masks = output0.length / 117;
const boxes = [];
for (let i = 0; i < num_masks; i++) {
const offset = i * 117;
const cls_conf = Array.from({ length: 80 }, (_, j) => output0[offset + 4 + j]);
const [classId, score] = cls_conf
.map((val, idx) => [idx, val])
.reduce((a, b) => (b[1] > a[1] ? b : a));
if (score < 0.5) continue;
const x = output0[offset];
const y = output0[offset + 1];
const w = output0[offset + 2];
const h = output0[offset + 3];
const mask_weights = output0.slice(offset + 4 + 80, offset + 4 + 80 + 32);
boxes.push({ x, y, w, h, score, classId, mask_weights });
}
const protoArr = new Float32Array(proto); // shape [1, 32, 160, 160]
const masks = [];
boxes.forEach((box) => {
const mask = new Float32Array(160 * 160).fill(0);
for (let i = 0; i < 32; i++) {
const weight = box.mask_weights[i];
for (let j = 0; j < 160 * 160; j++) {
mask[j] += protoArr[i * 160 * 160 + j] * weight;
}
}
// Apply sigmoid and resize to image size
const binaryMask = new Uint8ClampedArray(img.width * img.height * 4).fill(0);
const scaleX = img.width / 160;
const scaleY = img.height / 160;
for (let y = 0; y < 160; y++) {
for (let x = 0; x < 160; x++) {
const val = 1 / (1 + Math.exp(-mask[y * 160 + x]));
if (val > 0.5) {
const px = Math.floor(x * scaleX);
const py = Math.floor(y * scaleY);
const index = (py * img.width + px) * 4;
binaryMask[index] = 0; // R
binaryMask[index + 1] = 255; // G
binaryMask[index + 2] = 0; // B
binaryMask[index + 3] = 150; // Alpha
}
}
}
masks.push(new ImageData(binaryMask, img.width, img.height));
});
// Draw image and masks
ctx.drawImage(img, 0, 0);
masks.forEach((mask) => {
ctx.putImageData(mask, 0, 0);
});
}
</script>
</body>
</html>
0 条评论