44 */
55const fs = require ( "fs" ) ;
66const path = require ( "path" ) ;
7- const { cv, drawBlueRect , runVideoDetection } = require ( "./utils" ) ;
7+ const { cv, runVideoDetection } = require ( "./utils" ) ;
88
99if ( ! cv . xmodules . dnn ) {
1010 throw new Error ( "exiting: opencv4nodejs compiled without dnn module" ) ;
@@ -41,6 +41,13 @@ const labels = fs
4141
4242// initialize tensorflow darknet model from modelFile
4343const net = cv . readNetFromDarknet ( cfgFile , weightsFile ) ;
44+ const allLayerNames = net . getLayerNames ( ) ;
45+ const unconnectedOutLayers = net . getUnconnectedOutLayers ( ) ;
46+
47+ // determine only the *output* layer names that we need from YOLO
48+ const layerNames = unconnectedOutLayers . map ( layerIndex => {
49+ return allLayerNames [ layerIndex - 1 ] ;
50+ } ) ;
4451
4552const classifyImg = img => {
4653 // object detection model works with 416 x 416 images
@@ -49,12 +56,9 @@ const classifyImg = img => {
4956 const [ imgHeight , imgWidth ] = img . sizes ;
5057
5158 // network accepts blobs as input
52- const inputBlob = cv . blobFromImage ( img , 1 / 255.0 , size , vec3 , true , true ) ;
59+ const inputBlob = cv . blobFromImage ( img , 1 / 255.0 , size , vec3 , true , false ) ;
5360 net . setInput ( inputBlob ) ;
5461
55- // specify two layers "yolo_16" and "yolo_23"
56- const layerNames = [ "yolo_16" , "yolo_23" ] ;
57-
5862 console . time ( "net.forward" ) ;
5963 // forward pass input through entire network
6064 const layerOutputs = net . forward ( layerNames ) ;
@@ -95,20 +99,25 @@ const classifyImg = img => {
9599
96100 indices . forEach ( i => {
97101 const rect = boxes [ i ] ;
98- const imgRect = new cv . Rect ( rect . x , rect . y , rect . width , rect . height ) ;
99- drawBlueRect ( img , imgRect ) ;
102+
103+ const pt1 = new cv . Point ( rect . x , rect . y ) ;
104+ const pt2 = new cv . Point ( rect . x + rect . width , rect . y + rect . height ) ;
105+ const rectColor = new cv . Vec ( 255 , 0 , 0 ) ;
106+ const rectThickness = 2 ;
107+ const rectLineType = cv . LINE_8 ;
108+
109+ // draw the rect for the object
110+ img . drawRectangle ( pt1 , pt2 , rectColor , rectThickness , rectLineType ) ;
111+
100112 const text = labels [ classIDs [ i ] ] ;
101- img . putText (
102- text ,
103- new cv . Point ( rect . x , rect . y + 0.1 * imgHeight ) ,
104- cv . FONT_ITALIC ,
105- 2 ,
106- {
107- color : new cv . Vec ( 255 , 0 , 0 ) ,
108- thickness : 2
109- }
110- ) ;
111- drawBlueRect ( img , imgRect ) ;
113+ const org = new cv . Point ( rect . x , rect . y + 15 ) ;
114+ const fontFace = cv . FONT_HERSHEY_SIMPLEX ;
115+ const fontScale = 0.5 ;
116+ const textColor = new cv . Vec ( 123 , 123 , 255 ) ;
117+ const thickness = 2 ;
118+
119+ // put text on the object
120+ img . putText ( text , org , fontFace , fontScale , textColor , thickness ) ;
112121 } ) ;
113122 }
114123 } ) ;
0 commit comments