117 lines
4.3 KiB
Python
117 lines
4.3 KiB
Python
import cv2
|
|
import numpy as np
|
|
from flask import Flask, Response, request
|
|
import json
|
|
import time
|
|
import threading
|
|
import sys
|
|
import os
|
|
import uuid
|
|
import paho.mqtt.client as mqtt
|
|
|
|
# Configurações iniciais
|
|
mqtt_client = mqtt.Client(f"client_weed_detector_{uuid.uuid4()}")
|
|
mqtt_client.connect("localhost", port=1883)
|
|
|
|
model_config = 'C:\\train\\models\\ervas\\ervas.cfg'
|
|
model_weights = 'C:\\train\\models\\ervas\\backup\\ervas_final.weights'
|
|
labels_path = 'C:\\train\\models\\ervas\\labels.txt'
|
|
|
|
with open(labels_path, 'rt') as f:
|
|
classes = f.read().rstrip('\n').split('\n')
|
|
|
|
# Utilizando a GPU com CUDA
|
|
net = cv2.dnn.readNetFromDarknet(model_config, model_weights)
|
|
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
|
|
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
|
|
|
|
# Configurações adicionais
|
|
output_folder = 'Python/Output/'
|
|
max_readings = int(sys.argv[1])
|
|
|
|
app = Flask(__name__)
|
|
|
|
def detect_objects(conf_threshold, nms_threshold, _camera_index):
|
|
cap = cv2.VideoCapture(_camera_index, cv2.CAP_DSHOW)
|
|
width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
|
|
height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
|
|
readings = []
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
blob = cv2.dnn.blobFromImage(frame, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
|
|
net.setInput(blob)
|
|
outs = net.forward(net.getUnconnectedOutLayersNames())
|
|
|
|
current_readings = process_detections(outs, width, height, conf_threshold, frame) # Passando 'frame' como argumento
|
|
|
|
global json_data
|
|
timestamp = time.time()
|
|
json_data = {'timestamp': timestamp, 'x_max': width, 'y_max': height, 'objetos': current_readings}
|
|
readings.append(json_data)
|
|
|
|
mqtt_client.publish(sys.argv[6], json.dumps(json_data).encode('utf-8'))
|
|
|
|
manage_output(readings, frame, max_readings)
|
|
|
|
ret, buffer = cv2.imencode('.jpg', frame)
|
|
frame = buffer.tobytes()
|
|
yield (b'--frame\r\n'
|
|
b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
|
|
|
|
def process_detections(outs, width, height, conf_threshold, frame): # Recebendo 'frame' como argumento
|
|
current_readings = []
|
|
for out in outs:
|
|
for detection in out:
|
|
scores = detection[5:]
|
|
class_id = np.argmax(scores)
|
|
confidence = scores[class_id]
|
|
if confidence > conf_threshold:
|
|
x, y, w, h = calculate_box(detection, width, height)
|
|
detection_info = create_detection_info(class_id, x, y, w, h, confidence)
|
|
current_readings.append(detection_info)
|
|
if sys.argv[5] == "1":
|
|
draw_predictions(frame, class_id, confidence, x, y, w, h) # 'frame' já está disponível aqui
|
|
return current_readings
|
|
|
|
def draw_predictions(frame, class_id, confidence, x, y, w, h): # 'frame' é utilizado aqui
|
|
cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
|
|
cv2.putText(frame, f'{classes[class_id]} {confidence:.2f}', (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
|
|
|
|
def calculate_box(detection, width, height):
|
|
center_x, center_y, w, h = detection[0:4] * np.array([width, height, width, height])
|
|
x, y = int(center_x - w / 2), int(center_y - h / 2)
|
|
return x, y, int(w), int(h)
|
|
|
|
def create_detection_info(class_id, x, y, w, h, confidence):
|
|
return {
|
|
'id': int(class_id),
|
|
'descricao': classes[class_id],
|
|
'x': x,
|
|
'y': y,
|
|
'largura': w,
|
|
'altura': h,
|
|
'confianca': float(confidence)
|
|
}
|
|
|
|
def manage_output(readings, frame, max_readings):
|
|
if len(readings) > max_readings:
|
|
readings.pop(0)
|
|
if not os.path.exists(output_folder):
|
|
os.makedirs(output_folder)
|
|
with open(output_folder + sys.argv[4], 'w') as file:
|
|
json.dump(readings, file, indent=4)
|
|
|
|
@app.route('/' + sys.argv[3], methods=['GET'])
|
|
def video_feed():
|
|
conf_threshold = float(request.args.get('conf_threshold'))
|
|
nms_threshold = float(request.args.get('nms_threshold'))
|
|
camera_index = int(request.args.get('camera_index'))
|
|
return Response(detect_objects(conf_threshold, nms_threshold, camera_index), mimetype='multipart/x-mixed-replace; boundary=frame')
|
|
|
|
if __name__ == '__main__':
|
|
threading.Thread(target=lambda: mqtt_client.publish(sys.argv[6], "OK")).start()
|
|
app.run(host='0.0.0.0', port=sys.argv[2], threaded=True, debug=False) |