File size: 3,759 Bytes
ed9f15f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
/**
* YOLO Postprocessing - Process YOLO detection outputs
*/
import { DETECTION_CONFIG } from '../config.js';
import { cropImage } from '../utils/imageUtils.js';
/**
* Postprocess YOLO detections
*/
export async function postprocessYOLO(results, originalImage) {
const output = results.output0 || results[Object.keys(results)[0]];
const data = output.data;
const dims = output.dims; // [1, 5, 8400]
const detections = [];
const { confThreshold, iouThreshold } = DETECTION_CONFIG.yolo;
// YOLOv8 output format: [batch, [x, y, w, h, conf], num_anchors]
// dims = [1, 5, 8400]
const numAnchors = dims[2]; // 8400
// Extract boxes with confidence above threshold
const boxes = [];
for (let i = 0; i < numAnchors; i++) {
// Data is organized as: [x_center, y_center, width, height, confidence] for each anchor
const x_center = data[i];
const y_center = data[numAnchors + i];
const width = data[2 * numAnchors + i];
const height = data[3 * numAnchors + i];
const confidence = data[4 * numAnchors + i];
if (confidence > confThreshold) {
// Convert from center format to corner format
const x1 = (x_center - width / 2) / 640 * originalImage.width;
const y1 = (y_center - height / 2) / 640 * originalImage.height;
const x2 = (x_center + width / 2) / 640 * originalImage.width;
const y2 = (y_center + height / 2) / 640 * originalImage.height;
boxes.push({
x1: Math.max(0, x1),
y1: Math.max(0, y1),
x2: Math.min(originalImage.width, x2),
y2: Math.min(originalImage.height, y2),
confidence: confidence
});
}
}
// Apply Non-Maximum Suppression
const selectedBoxes = nonMaximumSuppression(boxes, iouThreshold);
// Crop each detected embryo
for (const box of selectedBoxes) {
const croppedData = await cropImage(originalImage, box.x1, box.y1, box.x2, box.y2);
detections.push({
box: box,
confidence: box.confidence,
imageData: croppedData
});
}
return detections;
}
/**
* Non-Maximum Suppression to remove overlapping boxes
*/
function nonMaximumSuppression(boxes, iouThreshold) {
// Sort boxes by confidence (descending)
boxes.sort((a, b) => b.confidence - a.confidence);
const selected = [];
const suppressed = new Set();
for (let i = 0; i < boxes.length; i++) {
if (suppressed.has(i)) continue;
selected.push(boxes[i]);
// Suppress overlapping boxes
for (let j = i + 1; j < boxes.length; j++) {
if (suppressed.has(j)) continue;
const iou = calculateIoU(boxes[i], boxes[j]);
if (iou > iouThreshold) {
suppressed.add(j);
}
}
}
return selected;
}
/**
* Calculate Intersection over Union (IoU) between two boxes
*/
function calculateIoU(box1, box2) {
const x1 = Math.max(box1.x1, box2.x1);
const y1 = Math.max(box1.y1, box2.y1);
const x2 = Math.min(box1.x2, box2.x2);
const y2 = Math.min(box1.y2, box2.y2);
const intersectionArea = Math.max(0, x2 - x1) * Math.max(0, y2 - y1);
const box1Area = (box1.x2 - box1.x1) * (box1.y2 - box1.y1);
const box2Area = (box2.x2 - box2.x1) * (box2.y2 - box2.y1);
const unionArea = box1Area + box2Area - intersectionArea;
return intersectionArea / unionArea;
}
|