136 lines
4.7 KiB
Python
136 lines
4.7 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
|
|
mqtt_client = mqtt.Client(f"client_weed_detector_{uuid.uuid4()}")
|
|
mqtt_client.connect("localhost", port=1883)
|
|
|
|
# Adiciona as configurações do modelo YOLO
|
|
model_config = 'C:\\ZendionINC\\agrobot_base\\Treinamento\\models\\ervas\\ervas.cfg'
|
|
model_weights = 'C:\\ZendionINC\\agrobot_base\\Treinamento\\models\\ervas\\backup\\ervas_final.weights'
|
|
labels_path = 'C:\\ZendionINC\\agrobot_base\\Treinamento\\models\\ervas\\labels.txt'
|
|
|
|
# Carregar as classes
|
|
with open(labels_path, 'rt') as f:
|
|
classes = f.read().rstrip('\n').split('\n')
|
|
|
|
# Carregar o modelo YOLO
|
|
net = cv2.dnn.readNetFromDarknet(model_config, model_weights)
|
|
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
|
|
net.setPreferableTarget(cv2.dnn.DNN_TARGET_OPENCL)
|
|
|
|
# Outras configurações
|
|
json_data = None
|
|
output_folder = 'Python/Output/'
|
|
max_readings = int(sys.argv[1])
|
|
porta = sys.argv[2]
|
|
url = sys.argv[3]
|
|
arquivoSaida = sys.argv[4]
|
|
mostrar_linhas = sys.argv[5] == "1"
|
|
mqtt_topic = sys.argv[6]
|
|
|
|
app = Flask(__name__)
|
|
|
|
# Função adaptada para detecção de ervas usando YOLO
|
|
def detect_objects(conf_threshold, nms_threshold, _camera_index):
|
|
#cap = cv2.VideoCapture(_camera_index, cv2.CAP_DSHOW)
|
|
cap = cv2.VideoCapture(_camera_index, cv2.CAP_MSMF)
|
|
width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
|
|
height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
|
|
readings = []
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
#frame = cv2.flip(frame, 1)
|
|
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 = []
|
|
for out in outs:
|
|
for detection in out:
|
|
scores = detection[5:]
|
|
class_id = np.argmax(scores)
|
|
confidence = scores[class_id]
|
|
if confidence > conf_threshold:
|
|
center_x = int(detection[0] * width)
|
|
center_y = int(detection[1] * height)
|
|
w = int(detection[2] * width)
|
|
h = int(detection[3] * height)
|
|
x = int(center_x - w / 2)
|
|
y = int(center_y - h / 2)
|
|
|
|
# Salvar as informações da detecção
|
|
detection_info = {
|
|
'id': int(class_id),
|
|
'descricao': classes[class_id],
|
|
'x': int(x),
|
|
'y': int(y),
|
|
'largura': int(w),
|
|
'altura': int(h),
|
|
'confianca': float(confidence)
|
|
}
|
|
current_readings.append(detection_info)
|
|
|
|
if mostrar_linhas:
|
|
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)
|
|
|
|
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(mqtt_topic, json.dumps(json_data).encode('utf-8'))
|
|
|
|
if not os.path.exists(output_folder):
|
|
os.makedirs(output_folder)
|
|
|
|
if len(readings) > max_readings:
|
|
readings.pop(0)
|
|
|
|
with open(output_folder + arquivoSaida, 'w') as file:
|
|
json.dump(readings, file, indent=4)
|
|
|
|
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')
|
|
|
|
|
|
|
|
# Esta função é para confirmar que script está pronto para execução
|
|
def send_script_ready():
|
|
mqtt_client.publish(mqtt_topic, "OK")
|
|
|
|
# Iniciar o servidor Flask em uma thread separada
|
|
def run_flask_server():
|
|
app.run(host='0.0.0.0', port=porta, threaded=True, debug=False)
|
|
|
|
@app.route('/' + url, 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')
|
|
|
|
|
|
# Iniciar o servidor
|
|
if __name__ == '__main__':
|
|
# Cria trhead separada para informar que o script iniciou com sucesso
|
|
mqtt_thread = threading.Thread(target=send_script_ready)
|
|
mqtt_thread.start()
|
|
|
|
# Inicia o servidor Flask na thread principal
|
|
run_flask_server()
|