Shinobi/plugins/python-yolo/pumpkin.py

94 lines
2.9 KiB
Python

from flask import Flask, request, jsonify
from pydarknet import Detector, Image
import cv2
import os
import json
import numpy as np
try:
with open('conf.json') as json_file:
config = json.load(json_file)
httpPort = config['pythonPort']
try:
httpPort
except NameError:
httpPort = 7990
except Exception as e:
print("conf.json not found.")
httpPort = 7990
# Load Flask
app = Flask("Pumpkin")
# 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"))
def spark(filepath):
try:
filepath
except NameError:
return "File path not found."
img = cv2.imread(filepath)
img2 = Image(img)
# r = net.classify(img2)
results = net.detect(img2)
returnData = '[]'
try:
print(type(results))
new_list = []
for item in results:
sub_list = {}
i = 0
for sub_item in item:
if i == 0:
key = 'tag'
sub_list[key] = sub_item.decode('utf-8')
if i == 1:
key = 'confidence'
sub_list[key] = sub_item
if i == 2:
key = 'points'
points_list = []
for points_item in sub_item:
points_list.append(points_item)
sub_list[key] = points_list
i += 1
new_list.append(sub_list)
returnData = jsonify(new_list)
# returnData = json.dumps(results)
except Exception as e:
returnData = ',\n'.join(map(str, results))
return returnData
# bake the image data by a file path
# POST body contains the "img" variable. The value should be to a local image path.
# Example : /dev/shm/streams/[GROUP_KEY]/[MONITOR_ID]/s.jpg
@app.route('/', methods=['GET'])
def index():
return "Pumpkin.py is running. This web interface should NEVER be accessible remotely. The Node.js plugin that runs this script should only be allowed accessible remotely."
# bake the image data by a file path
# POST body contains the "img" variable. The value should be to a local image path.
# Example : /dev/shm/streams/[GROUP_KEY]/[MONITOR_ID]/s.jpg
@app.route('/post', methods=['POST'])
def post():
filepath = request.form['img']
return 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.
# Example : /dev/shm/streams/[GROUP_KEY]/[MONITOR_ID]/s.jpg
@app.route('/get', methods=['GET'])
def get():
filepath = request.args.get('img')
return spark(filepath)
# quick-and-dirty start
if __name__ == '__main__':
app.run(host='0.0.0.0', port=httpPort)