From b91f64c0fc70547a362d3732deb0bfd506979ca9 Mon Sep 17 00:00:00 2001 From: Moe Date: Sat, 11 Aug 2018 11:54:41 -0700 Subject: [PATCH] update Python-Yolo plugin to use socket.io instead of http for detections much faster now :D --- plugins/python-yolo/INSTALL.sh | 10 ++---- plugins/python-yolo/bootPy.sh | 1 + plugins/python-yolo/package.json | 8 ++--- plugins/python-yolo/pumpkin.py | 27 +++++++++----- plugins/python-yolo/shinobi-python-yolo.js | 42 ++++++++++++++++++---- 5 files changed, 62 insertions(+), 26 deletions(-) create mode 100644 plugins/python-yolo/bootPy.sh diff --git a/plugins/python-yolo/INSTALL.sh b/plugins/python-yolo/INSTALL.sh index ef32c6d4..a7d1f80b 100644 --- a/plugins/python-yolo/INSTALL.sh +++ b/plugins/python-yolo/INSTALL.sh @@ -37,13 +37,9 @@ fi echo "-----------------------------------" sudo apt update -y sudo apt-get install libxml2-dev libxslt-dev libxslt1-dev zlib1g-dev -y -if ! [ -x "$(command -v python3)" ]; then - echo "Installing python3" - sudo apt install python3 python3-dev python3-pip -y - sudo apt install python3-lxml libxml2-dev -y -else - echo "python3 found... : $(python3 --version)" -fi +echo "Installing python3" +sudo apt install python3 python3-dev python3-pip -y +sudo apt install python3-lxml libxml2-dev -y echo "-----------------------------------" if ! [ -x "$(command -v gcc-6)" ]; then echo "Installing gcc-6 and g++-6" diff --git a/plugins/python-yolo/bootPy.sh b/plugins/python-yolo/bootPy.sh new file mode 100644 index 00000000..68be5794 --- /dev/null +++ b/plugins/python-yolo/bootPy.sh @@ -0,0 +1 @@ +python3 -u $@ \ No newline at end of file diff --git a/plugins/python-yolo/package.json b/plugins/python-yolo/package.json index 7f6d5aa0..544f33b8 100644 --- a/plugins/python-yolo/package.json +++ b/plugins/python-yolo/package.json @@ -1,10 +1,10 @@ { - "name": "shinobi-openalpr", + "name": "shinobi-python-yolo", "version": "1.0.0", - "description": "OpenALPR plugin for Shinobi", - "main": "shinobi-openalpr.js", + "description": "YOLOv3 plugin for Shinobi that uses Python functions for detection.", + "main": "shinobi-python-yolo.js", "dependencies": { - "canvas": "^1.6.7", + "socket.io-client": "^1.7.4", "express": "^4.16.2", "moment": "^2.19.2", "socket.io": "^2.0.4" diff --git a/plugins/python-yolo/pumpkin.py b/plugins/python-yolo/pumpkin.py index eb6e84c0..c26a2d53 100644 --- a/plugins/python-yolo/pumpkin.py +++ b/plugins/python-yolo/pumpkin.py @@ -1,11 +1,16 @@ -from flask import Flask, request, jsonify +from flask import Flask, request, jsonify, render_template +from flask_socketio import SocketIO, emit from pydarknet import Detector, Image import cv2 import os import json import numpy as np +import sys + +dirname = sys.argv[1] + try: - with open('conf.json') as json_file: + with open("{}/conf.json".format(dirname)) as json_file: config = json.load(json_file) httpPort = config['pythonPort'] try: @@ -17,14 +22,15 @@ except Exception as e: httpPort = 7990 # Load Flask -app = Flask("Pumpkin") +app = Flask("YOLOv3 for Shinobi (Pumpkin Pie)") +socketio = SocketIO(app) # Silence Flask import logging log = logging.getLogger('werkzeug') log.setLevel(logging.ERROR) # Load Darknet -net = Detector(bytes("cfg/yolov3.cfg", encoding="utf-8"), bytes("weights/yolov3.weights", encoding="utf-8"), 0, bytes("cfg/coco.data",encoding="utf-8")) +net = Detector(bytes("{}/cfg/yolov3.cfg".format(dirname), encoding="utf-8"), bytes("{}/weights/yolov3.weights".format(dirname), encoding="utf-8"), 0, bytes("{}/cfg/coco.data".format(dirname),encoding="utf-8")) def spark(filepath): try: @@ -39,7 +45,6 @@ def spark(filepath): results = net.detect(img2) returnData = '[]' try: - print(type(results)) new_list = [] for item in results: sub_list = {} @@ -59,7 +64,7 @@ def spark(filepath): sub_list[key] = points_list i += 1 new_list.append(sub_list) - returnData = jsonify(new_list) + returnData = new_list # returnData = json.dumps(results) except Exception as e: returnData = ',\n'.join(map(str, results)) @@ -79,7 +84,7 @@ def index(): @app.route('/post', methods=['POST']) def post(): filepath = request.form['img'] - return spark(filepath) + return jsonify(spark(filepath)) # bake the image data by a file path # GET string contains the "img" variable. The value should be to a local image path. @@ -87,8 +92,12 @@ def post(): @app.route('/get', methods=['GET']) def get(): filepath = request.args.get('img') - return spark(filepath) + return jsonify(spark(filepath)) + +@socketio.on('f') +def receiveMessage(message): + emit('f',{'id':message.get("id"),'data':spark(message.get("path"))}) # quick-and-dirty start if __name__ == '__main__': - app.run(host='0.0.0.0', port=httpPort) \ No newline at end of file + socketio.run(app, port=httpPort) \ No newline at end of file diff --git a/plugins/python-yolo/shinobi-python-yolo.js b/plugins/python-yolo/shinobi-python-yolo.js index c223b5d0..998fd2d3 100644 --- a/plugins/python-yolo/shinobi-python-yolo.js +++ b/plugins/python-yolo/shinobi-python-yolo.js @@ -17,7 +17,8 @@ var spawn = require('child_process').spawn; var moment = require('moment'); var http = require('http'); var express = require('express'); -var config=require('./conf.json'); +var socketIoClient = require('socket.io-client'); +var config = require('./conf.json'); var http = require('http'), app = express(), server = http.createServer(app); @@ -78,6 +79,8 @@ s.getRequest = function(url,callback){ // s.systemLog("Get Snapshot Error", e); }); } +s.multiplerHeight = 0.75 +s.multiplerWidth = 0.96 s.detectObject=function(buffer,d,tx){ d.tmpFile=s.gid(5)+'.jpg' if(!fs.existsSync(s.dir.streams)){ @@ -94,15 +97,17 @@ s.detectObject=function(buffer,d,tx){ fs.writeFile(d.dir+d.tmpFile,buffer,function(err){ if(err) return s.systemLog(err); if(s.isPythonRunning === false)return console.log('Python Script is not Running.') - s.getRequest('http://localhost:'+config.pythonPort+'/get?img='+d.dir+d.tmpFile,function(data){ + var callbackId = s.gid(10) + s.sendToPython({path:d.dir+d.tmpFile,id:callbackId},function(data){ if(data.length > 0){ var mats=[] data.forEach(function(v){ mats.push({ - x:v.points[0], - y:v.points[1], + x:v.points[0] * s.multiplerWidth, + y:v.points[1] * s.multiplerHeight, width:v.points[2], height:v.points[3], + confidence:v.confidence, tag:v.tag }) }) @@ -120,6 +125,7 @@ s.detectObject=function(buffer,d,tx){ } }) } + delete(s.callbacks[callbackId]) exec('rm -rf '+d.dir+d.tmpFile,{encoding:'utf8'}) }) }) @@ -207,7 +213,7 @@ if(config.mode==='host'){ }else{ //start plugin as client if(!config.host){config.host='localhost'} - var io = require('socket.io-client')('ws://'+config.host+':'+config.port);//connect to master + var io = socketIoClient('ws://'+config.host+':'+config.port);//connect to master s.cx=function(x){x.pluginKey=config.key;x.plug=config.plug;return io.emit('ocv',x)} io.on('connect',function(d){ s.cx({f:'init',plug:config.plug,notice:config.notice,type:config.type}); @@ -220,10 +226,34 @@ if(config.mode==='host'){ }) } +//Start Python Controller +var pythonIo = socketIoClient('ws://localhost:'+config.pythonPort); +s.callbacks = {} +s.sendToPython = function(data,callback){ + s.callbacks[data.id] = callback + pythonIo.emit('f',data) +} +pythonIo.on('connect',function(d){ + console.log('Connected') +}) +pythonIo.on('disconnect',function(d){ + setTimeout(function(){ + pythonIo.connect(); + console.log('Disconnected') + },3000) +}) +pythonIo.on('f',function(d){ + if(s.callbacks[d.id]){ + s.callbacks[d.id](d.data) + } +}) + + //Start Python Daemon +process.env.PYTHONUNBUFFERED = 1; s.createPythonProcess = function(){ s.isPythonRunning = false - s.pythonScript = spawn('python3',[config.pythonScript]); + s.pythonScript = spawn('sh',[__dirname+'/bootPy.sh',config.pythonScript,__dirname]); var onStdErr = function (data) { console.log('Python ERR') console.log(data.toString())