// Reference the elements that we will need
const oContainer = document.getElementById('container');
const oOverlay = document.getElementById('overlay');
const oCanvas = document.getElementById('canvas');
const oVideo = document.getElementById('video');
// if we detect only 1 class, we technically only need 1 colour
const COLOURS = [
"#EF4444", "#4299E1", "#059669",
"#FBBF24", "#4B52B1", "#7B3AC2",
"#ED507A", "#1DD1A1", "#F3873A",
"#4B5563", "#DC2626", "#1852B4",
"#18A35D", "#F59E0B", "#4059BE",
"#6027A5", "#D63D60", "#00AC9B",
"#E64A19", "#272A34"
];
let bIsProcessing = false;
const oContext = oCanvas.getContext('2d', {
willReadFrequently: true
});
const sModelId = 'data'; // this is the assets folder where both .config and .onnx are stored
let oModel, oProcessor;
let iThreshold = 0.25;
(async () => {
try {
// init transformers.js environment variables
transformers.env.allowRemoteModels = false;
transformers.env.allowLocalModels = true;
transformers.env.localModelPath = '.';
// Load the fine-tuned model
oModel = await transformers.AutoModel.from_pretrained(sModelId, {
model_file_name: 'YOLO_with_NMS', // adapt this to match your model name
subfolder: '',
dtype: 'fp32', // new transformers.js version (after 3.x) full-precision model: 'fp32', quantized model: 'int8'
});
sModelInputName = oModel.sessions.model.inputNames[0];
sModelOutputName = oModel.sessions.model.outputNames[0];
// Load the image processor
oProcessor = await transformers.AutoProcessor.from_pretrained(sModelId);
// Start the video stream
navigator.mediaDevices.getUserMedia({
video: {
facingMode: 'environment'
}
}, // Ask for video
).then((stream) => {
// Set up the video and canvas elements.
oVideo.srcObject = oStream;
oVideo.play();
const oVideoTrack = oStream.getVideoTracks()[0];
const {
width,
height
} = oVideoTrack.getSettings();
oCanvas.width = width;
oCanvas.height = height;
// Set container width and height depending on the image aspect ratio
const fAspectRatio = width / height;
const [iContainerWidth, iContainerHeight] = (fAspectRatio > 720 / 405) ? [720, 720 / fAspectRatio] : [405 * fAspectRatio, 405];
oContainer.style.width = `${iContainerWidth}px`;
oContainer.style.height = `${iContainerHeight}px`;
// Start the animation loop
window.requestAnimationFrame(updateCanvas);
fnSuccess({});
});
} catch (e) {
fnError(e.message);
}
})();
// Run processor and model on pixel data every frame
function updateCanvas() {
const {
width,
height
} = oCanvas;
oContext.drawImage(oVideo, 0, 0, width, height);
if (!bIsProcessing) {
bIsProcessing = true;
(async function() {
try {
// Read the current frame from the video
const oPixelData = oContext.getImageData(0, 0, width, height).data;
const oRawImage = new transformers.RawImage(oPixelData, width, height, 4);
// Process the image
const {
pixel_values,
reshaped_input_sizes
} = await oProcessor(oRawImage);
// Run the model
const aModelOutputs = await oModel({
[sModelInputName]: pixel_values
});
// Update UI
oOverlay.innerHTML = '';
// Draw boxes on video
const aSizes = reshaped_input_sizes[0].reverse();
// Iterate over all detected boxes and check the box data
aModelOutputs[sModelOutputName].tolist()[0].forEach(aBox => renderBox(aBox, aSizes));
bIsProcessing = false;
} catch (e) {
fnError(e.message);
}
})();
}
window.requestAnimationFrame(updateCanvas);
}
// Render a bounding box and label on the image
function renderBox([iXmin, iYmin, iXmax, iYmax, fScore, iClassId], [iImageWidth, iImageHeight]) {
try {
if (fScore < fConfidenceThreshold) return; // Skip boxes with low confidence
// Generate a random color for the box
const sColour = COLOURS[iClassId % COLOURS.length];
// Draw the box
const oBoxElement = document.createElement('div');
oBoxElement.className = 'bounding-box';
Object.assign(oBoxElement.style, {
borderColor: sColour,
left: 100 * iXmin / iImageWidth + '%',
top: 100 * iYmin / iImageHeight + '%',
width: 100 * (iXmax - iXmin) / iImageWidth + '%',
height: 100 * (iYmax - iYmin) / iImageHeight + '%',
})
// Draw label
const oLabelElement = document.createElement('span');
oLabelElement.textContent = `${oModel.config.id2label[iClassId]} (${(100 * fScore).toFixed(2)}%)`;
oLabelElement.className = 'bounding-box-label';
oLabelElement.style.backgroundColor = sColour;
oBoxElement.appendChild(oLabelElement);
oOverlay.appendChild(oBoxElement);
} catch (e) {
fnError(e.message);
}
}
