amas_ai_webmodel
Version:
YOLO Inference library using TensorFlow.js
3 lines (2 loc) • 4.82 kB
JavaScript
import*as e from"@tensorflow/tfjs";import"@tensorflow/tfjs-backend-webgl";function t(e,t,o,i){return new(o||(o=Promise))((function(s,n){function r(e){try{l(i.next(e))}catch(e){n(e)}}function a(e){try{l(i.throw(e))}catch(e){n(e)}}function l(e){var t;e.done?s(e.value):(t=e.value,t instanceof o?t:new o((function(e){e(t)}))).then(r,a)}l((i=i.apply(e,t||[])).next())}))}"function"==typeof SuppressedError&&SuppressedError;class o{constructor(){this.config={modelUrl:"",labels:["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"],colors:["#FF6B6B","#FFA372","#FFDC64","#A3E635","#66D3FA","#82AAFF","#C084FC","#F472B6","#FF9F43","#FFB74D","#AED581","#4DD0E1","#4FC3F7","#9575CD","#F06292","#BA68C8","#FF8A65","#FFD54F","#81C784","#64B5F6"],displayLabels:null,scoreThreshold:.5,boxLineWidth:2,boxLabels:!0}}setup(e){e.modelUrl&&(this.config.modelUrl=e.modelUrl),e.labels&&(this.config.labels=e.labels),e.colors&&(this.config.colors=e.colors),e.displayLabels&&(this.config.displayLabels=new Set(e.displayLabels)),void 0!==e.scoreThreshold&&(this.config.scoreThreshold=e.scoreThreshold),e.boxLineWidth&&(this.config.boxLineWidth=e.boxLineWidth),"boxLabels"in e&&(this.config.boxLabels=Boolean(e.boxLabels))}loadModel(){return t(this,void 0,void 0,(function*(){yield e.ready();try{const t=yield e.loadGraphModel(this.config.modelUrl),o=e.ones(t.inputs[0].shape),i=t.execute(o);return e.dispose([o,i]),{net:t,inputShape:t.inputs[0].shape}}catch(e){return console.error("Error loading model:",e),null}}))}detect(o,i,s){return t(this,arguments,void 0,(function*(t,o,i,s=()=>{}){const[n,r]=o.inputShape.slice(1,3);e.engine().startScope();const[a,{scale:l,dx:c,dy:d,origWidth:h,origHeight:f}]=this.preprocess(t,n,r),b=o.net.execute(a),p=b.transpose([0,2,1]),u=e.tidy((()=>{const t=p.slice([0,0,0],[-1,-1,1]),o=p.slice([0,0,1],[-1,-1,1]),i=p.slice([0,0,2],[-1,-1,1]),s=p.slice([0,0,3],[-1,-1,1]),n=e.sub(t,e.div(i,2)),r=e.sub(o,e.div(s,2)),a=e.add(t,e.div(i,2)),l=e.add(o,e.div(s,2));return e.concat([r,n,l,a],2).squeeze()})),[g,m]=e.tidy((()=>{const e=p.slice([0,0,4],[-1,-1,this.config.labels.length]).squeeze([0]);return[e.max(1),e.argMax(1)]})),y=yield e.image.nonMaxSuppressionAsync(u,g,500,.45,.2),x=u.gather(y,0).dataSync(),F=g.gather(y,0).dataSync(),v=m.gather(y,0).dataSync(),k=[];for(let e=0;e<F.length;e++){let[t,o,i,s]=x.slice(4*e,4*(e+1));o=Math.max(0,(o-c)/l),s=Math.min(h,(s-c)/l),t=Math.max(0,(t-d)/l),i=Math.min(f,(i-d)/l),k.push(t,o,i,s)}HTMLImageElement,i.width=h,i.height=f,this.renderBoxes(i,k,F,v,[1,1]);const w=[];for(let e=0;e<v.length;e++){const t=v[e],o=this.config.labels[t];w.push(o)}s({boxes:k,scores:Array.from(F),classes:Array.from(v),labels:w}),e.dispose([b,p,u,g,m,y]),e.engine().endScope()}))}detectVideo(e,o,i){const s=()=>t(this,void 0,void 0,(function*(){if(0===e.videoWidth&&null===e.srcObject){const e=i.getContext("2d");return e&&e.clearRect(0,0,i.width,i.height),void requestAnimationFrame(s)}yield this.detect(e,o,i,(()=>{requestAnimationFrame(s)}))}));s()}preprocess(t,o,i){const s=e.browser.fromPixels(t),[n,r]=s.shape.slice(0,2),a=Math.min(o/r,i/n),l=Math.round(r*a),c=Math.round(n*a),d=e.image.resizeBilinear(s,[c,l]).div(255),h=Math.floor((o-l)/2),f=Math.floor((i-c)/2),b=e.tidy((()=>e.pad(d,[[f,i-c-f],[h,o-l-h],[0,0]]))),p=b.expandDims(0);return e.dispose([s,d,b]),[p,{scale:a,dx:h,dy:f,origWidth:r,origHeight:n}]}renderBoxes(e,t,o,i,s){const n=e.getContext("2d");if(n){n.clearRect(0,0,e.width,e.height),n.font="16px Arial",n.textBaseline="top";for(let e=0;e<o.length;e++){if(o[e]<this.config.scoreThreshold)continue;const s=this.config.labels[i[e]];if(this.config.displayLabels&&!this.config.displayLabels.has(s))continue;const r=this.config.colors[i[e]%this.config.colors.length],a=(100*o[e]).toFixed(1);let[l,c,d,h]=t.slice(4*e,4*(e+1));const f=h-c,b=d-l;if(n.strokeStyle=r,n.lineWidth=this.config.boxLineWidth,n.strokeRect(c,l,f,b),this.config.boxLabels){const e=`${s} - ${a}%`,t=n.measureText(e).width,o=16;n.fillStyle=r,n.fillRect(c,l-o,t+4,o),n.fillStyle="#ffffff",n.fillText(e,c+2,l-o)}}}}}export{o as default};
//# sourceMappingURL=yolo.esm.js.map