ppu-yolo-onnx-inference/core/base-yolo-inference.js

Use your YOLO onnx object detection model in Typescript Bun environment easily.

export class BaseYoloDetectionInference{model;classNames;thresholds;debugging;platform;modelMetadata=null;session=null;static CHANNELS=3;constructor(options,platform){this.model=options.model.onnx;this.classNames=options.model.classNames;this.thresholds={...DEFAULT_THRESHOLDS,...options.thresholds};this.debugging={...DEFAULT_DEBUG_OPTIONS,...options.debug};this.modelMetadata=options.modelMetadata||null;this.platform=platform}isInitialized(){return this.session!==null}async init(){if(!this.model){throw new Error("Model ArrayBuffer is required")}try{this.log("init","Loading model");this.session=await this.platform.createInferenceSession(this.model);await this.platform.scheduleYield();this.log("init",`Model loaded successfully
	inputNames: ${this.session.inputNames}
	outputNames: ${this.session.outputNames}`);if(!this.modelMetadata){let inputShape=STANDARD_MODEL_INPUT_SHAPE;try{let meta=this.session.inputMetadata;if(Array.isArray(meta)&&meta.length>0){let shape=meta[0].shape.slice(2);if(shape.length===2)inputShape=shape}}catch{}this.modelMetadata={inputTensorName:this.session.inputNames[0],outputTensorName:this.session.outputNames[0],inputShape}}else{this.modelMetadata={inputTensorName:this.modelMetadata.inputTensorName||this.session.inputNames[0],outputTensorName:this.modelMetadata.outputTensorName||this.session.outputNames[0],inputShape:this.modelMetadata.inputShape||STANDARD_MODEL_INPUT_SHAPE}}await this.platform.initRuntime();this.log("init","ImageProcessor runtime initialized")}catch(error){throw new Error(`Failed to load model: ${error}`)}}async detect(image){this.log("detect","Starting object detection");try{let preprocessed=await this.preprocessImage(image);let output=await this.runInference(preprocessed);if(!output)return[];if(this.debugging.verbose){this.debugTensorData(output.data,output.dims[2],output.dims[1])}let detections=this.postprocessOutput(output,preprocessed);if(this.debugging.debug){await this.saveDebugImages(preprocessed,image,detections)}this.log("detect",`Detected ${detections.length} objects
`);return detections}catch(error){console.error("Detection failed:",error);return[]}}async preprocessImage(image){if(!this.modelMetadata){throw new Error("Model not initialized. Call init() first")}let canvas=this.platform.isCanvasLike(image)?image:await this.platform.prepareCanvas(image);const{width:originalWidth,height:originalHeight}=canvas;const[modelInputWidth,modelInputHeight]=this.modelMetadata.inputShape;let scaleRatio=Math.min(modelInputWidth/originalWidth,modelInputHeight/originalHeight);let scaledWidth=Math.round(originalWidth*scaleRatio);let scaledHeight=Math.round(originalHeight*scaleRatio);let padX=(modelInputWidth-scaledWidth)/2;let padY=(modelInputHeight-scaledHeight)/2;let scaledCanvas=this.platform.createCanvas(modelInputWidth,modelInputHeight);let ctx=this.platform.getContext2D(scaledCanvas);ctx.fillStyle="rgb(114, 114, 114)";ctx.fillRect(0,0,modelInputWidth,modelInputHeight);let resized=this.platform.resizeImage(canvas,scaledWidth,scaledHeight);ctx.drawImage(resized,padX,padY,scaledWidth,scaledHeight);let tensor=this.canvasToTensor(scaledCanvas,modelInputWidth,modelInputHeight);this.log("preprocessImage",`Preprocessed: ${originalWidth}x${originalHeight} → ${modelInputWidth}x${modelInputHeight} (ratio: ${scaleRatio.toFixed(3)}, padX: ${padX.toFixed(1)}, padY: ${padY.toFixed(1)})`);return{tensor,modelInputWidth,modelInputHeight,originalWidth,originalHeight,scaleRatio,padX,padY}}canvasToTensor(canvas,width,height){let tensor=new Float32Array(BaseYoloDetectionInference.CHANNELS*height*width);let ctx=this.platform.getContext2D(canvas);let imageData=ctx.getImageData(0,0,width,height).data;for(let h=0;h<height;h++){for(let w=0;w<width;w++){let pixelIndex=h*width+w;let rgbaIndex=pixelIndex*4;tensor[pixelIndex]=imageData[rgbaIndex]/255;tensor[height*width+pixelIndex]=imageData[rgbaIndex+1]/255;tensor[2*height*width+pixelIndex]=imageData[rgbaIndex+2]/255}}return tensor}async runInference(preprocessed){if(!this.session||!this.modelMetadata){throw new Error("Model not initialized. Call init() first")}try{let inputTensor=this.platform.createTensor("float32",preprocessed.tensor,[1,3,preprocessed.modelInputHeight,preprocessed.modelInputWidth]);let feeds={[this.modelMetadata.inputTensorName]:inputTensor};let results=await this.session.run(feeds);return results[this.modelMetadata.outputTensorName]||null}catch(error){console.error("Inference error:",error);throw error}}postprocessOutput(tensor,preprocessed){let data=tensor.data;const[,numParams,numPredictions]=tensor.dims;this.log("postprocessOutput",`Post-processing output: numParams=${numParams}, numPredictions=${numPredictions}`);if(numParams<4){console.error(`Invalid tensor shape: expected ≥4 parameters per box, got ${numParams}`);return[]}let candidates=this.extractCandidates(data,numPredictions,numParams,preprocessed);let nmsIndices=this.applyNMS(candidates);return this.scaleCandidates(candidates,nmsIndices,preprocessed)}extractCandidates(data,numPredictions,numParams,_preprocessed){let candidates=[];let numClasses=numParams-4;let isSingleClass=numClasses<=1;this.log("extractCandidates",`YOLOv11 format: numClasses=${numClasses}, isSingleClass=${isSingleClass}`);let debugCount=0;let highConfidenceCount=0;for(let i=0;i<numPredictions;i++){let cx=data[i];let cy=data[numPredictions+i];let w=data[2*numPredictions+i];let h=data[3*numPredictions+i];let finalConfidence;let bestClassId;if(isSingleClass){finalConfidence=data[4*numPredictions+i];bestClassId=0}else{let maxClassScore=-1/0;bestClassId=0;for(let c=0;c<numClasses;c++){let classScore=data[(4+c)*numPredictions+i];if(classScore>maxClassScore){maxClassScore=classScore;bestClassId=c}}finalConfidence=maxClassScore}if(finalConfidence>0.1)highConfidenceCount++;if(finalConfidence>0.1&&debugCount<5){this.log("extractCandidates",`Debug candidate ${i}: confidence=${finalConfidence.toFixed(4)}, `+`box=[${cx.toFixed(1)}, ${cy.toFixed(1)}, ${w.toFixed(1)}, ${h.toFixed(1)}], `+`classId=${bestClassId}`);debugCount++}if(finalConfidence<this.thresholds.confidence||w<=0||h<=0)continue;let x=cx-w/2;let y=cy-h/2;if(!this.modelMetadata)continue;if(x<0||y<0||x+w>this.modelMetadata.inputShape[0]||y+h>this.modelMetadata.inputShape[1]){continue}candidates.push({box:{x,y,width:w,height:h},score:finalConfidence,classId:bestClassId})}this.log("extractCandidates",`Total candidates with confidence > 0.1: ${highConfidenceCount}
[YOLO:extractCandidates] `+`Final candidates after filtering: ${candidates.length}`);if(candidates.length<3&&this.thresholds.confidence>0.1){this.log("extractCandidates","Few candidates found, trying with lower threshold...");return this.extractWithLowerThreshold(data,numPredictions,numParams)}return candidates}extractWithLowerThreshold(data,numPredictions,numParams){let candidates=[];let numClasses=numParams-4;let isSingleClass=numClasses<=1;let lowerThreshold=Math.max(0.05,this.thresholds.confidence*0.5);this.log("extractWithLowerThreshold",`Trying lower threshold: ${lowerThreshold}`);for(let i=0;i<numPredictions;i++){let cx=data[i];let cy=data[numPredictions+i];let w=data[2*numPredictions+i];let h=data[3*numPredictions+i];let finalConfidence;let bestClassId;if(isSingleClass){finalConfidence=data[4*numPredictions+i];bestClassId=0}else{let maxClassScore=-1/0;bestClassId=0;for(let c=0;c<numClasses;c++){let classScore=data[(4+c)*numPredictions+i];if(classScore>maxClassScore){maxClassScore=classScore;bestClassId=c}}finalConfidence=maxClassScore}if(finalConfidence>=lowerThreshold&&w>0&&h>0){let x=cx-w/2;let y=cy-h/2;if(!this.modelMetadata)continue;if(x>=0&&y>=0&&x+w<=this.modelMetadata.inputShape[0]&&y+h<=this.modelMetadata.inputShape[1]){candidates.push({box:{x,y,width:w,height:h},score:finalConfidence,classId:bestClassId})}}}this.log("extractWithLowerThreshold",`Candidates with lower threshold: ${candidates.length}`);return candidates}debugTensorData(data,numPredictions,numParams){this.log("debugTensorData",`
=== TENSOR DEBUG ===`);let confidences=[];for(let i=0;i<numPredictions;i++){if(numParams===5){let conf=data[4*numPredictions+i];confidences.push({index:i,confidence:conf})}else{let maxConf=-1/0;for(let c=0;c<numParams-4;c++){let classScore=data[(4+c)*numPredictions+i];maxConf=Math.max(maxConf,classScore)}confidences.push({index:i,confidence:maxConf})}}confidences.sort((a,b)=>b.confidence-a.confidence);this.log("debugTensorData","Top 10 detections by confidence:");for(let j=0;j<Math.min(10,confidences.length);j++){const{index:i,confidence}=confidences[j];let cx=data[i];let cy=data[numPredictions+i];let w=data[2*numPredictions+i];let h=data[3*numPredictions+i];this.log("debugTensorData",`${j+1}. Index ${i}: conf=${confidence.toFixed(4)}, box=[${cx.toFixed(1)}, ${cy.toFixed(1)}, ${w.toFixed(1)}, ${h.toFixed(1)}]`)}let ranges=[0,0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1];let counts=new Array(ranges.length-1).fill(0);confidences.forEach(({confidence})=>{for(let r=0;r<ranges.length-1;r++){if(confidence>=ranges[r]&&confidence<ranges[r+1]){counts[r]++;break}}});this.log("debugTensorData","Confidence distribution:");for(let r=0;r<counts.length;r++){this.log("debugTensorData",`${ranges[r].toFixed(1)}-${ranges[r+1].toFixed(1)}: ${counts[r]} detections`)}}applyNMS(candidates){let indices=candidates.map((_,i)=>i).sort((a,b)=>candidates[b].score-candidates[a].score);let keep=[];while(indices.length>0){let current=indices.shift();keep.push(current);for(let i=indices.length-1;i>=0;i--){let iou=this.calculateIoU(candidates[current].box,candidates[indices[i]].box);if(iou>this.thresholds.iou){indices.splice(i,1)}}}return keep}scaleCandidates(candidates,indices,preprocessed){return indices.map((i)=>{let candidate=candidates[i];const{scaleRatio,originalWidth,originalHeight,padX,padY}=preprocessed;let unpaddedX=candidate.box.x-padX;let unpaddedY=candidate.box.y-padY;let x=Math.max(0,Math.round(unpaddedX/scaleRatio));let y=Math.max(0,Math.round(unpaddedY/scaleRatio));let width=Math.min(originalWidth-x,Math.round(candidate.box.width/scaleRatio));let height=Math.min(originalHeight-y,Math.round(candidate.box.height/scaleRatio));if(width<=0||height<=0)return null;return{box:{x,y,width,height},className:this.classNames[candidate.classId]||`class_${candidate.classId}`,classId:candidate.classId,confidence:candidate.score}}).filter(Boolean)}calculateIoU(box1,box2){let x1=Math.max(box1.x,box2.x);let y1=Math.max(box1.y,box2.y);let x2=Math.min(box1.x+box1.width,box2.x+box2.width);let y2=Math.min(box1.y+box1.height,box2.y+box2.height);let intersection=Math.max(0,x2-x1)*Math.max(0,y2-y1);let union=box1.width*box1.height+box2.width*box2.height-intersection;return union>0?intersection/union:0}async saveDebugImages(preprocessed,originalImage,detections){try{await this.savePreprocessedImage(preprocessed);await this.saveDetectionVisualization(originalImage,detections)}catch(error){console.error("Debug image save failed:",error)}}async savePreprocessedImage(preprocessed){const{modelInputWidth,modelInputHeight,tensor}=preprocessed;let canvas=this.platform.createCanvas(modelInputWidth,modelInputHeight);let ctx=this.platform.getContext2D(canvas);let imageData=ctx.createImageData(modelInputWidth,modelInputHeight);for(let i=0;i<modelInputWidth*modelInputHeight;i++){let rgbaIndex=i*4;imageData.data[rgbaIndex]=tensor[i]*255;imageData.data[rgbaIndex+1]=tensor[modelInputWidth*modelInputHeight+i]*255;imageData.data[rgbaIndex+2]=tensor[2*modelInputWidth*modelInputHeight+i]*255;imageData.data[rgbaIndex+3]=255}ctx.putImageData(imageData,0,0);await this.platform.saveDebugImage(canvas,"yolo-preprocessed",this.debugging.debugFolder)}async saveDetectionVisualization(image,detections){let canvas=this.platform.isCanvasLike(image)?image:await this.platform.prepareCanvas(image);let ctx=this.platform.getContext2D(canvas);detections.forEach((detection)=>{const{x,y,width,height}=detection.box;ctx.strokeStyle="yellow";ctx.lineWidth=3;ctx.strokeRect(x,y,width,height);ctx.fillStyle="yellow";ctx.font="16px Arial";ctx.fillText(`${detection.className} ${(detection.confidence*100).toFixed(1)}%`,x,y>20?y-5:y+height+15)});await this.platform.saveDebugImage(canvas,"yolo-detections",this.debugging.debugFolder)}log(caller,message){if(this.debugging.verbose){console.log(`[YOLO:${caller}] ${message}`)}}async destroy(){this.log("destroy","Cleaning up resources");if(this.session){await this.session.release();this.session=null}this.log("destroy","Resources cleaned up")}}import{DEFAULT_DEBUG_OPTIONS,DEFAULT_THRESHOLDS,STANDARD_MODEL_INPUT_SHAPE}from"../constant.js";