145 lines
5.1 KiB
Python
145 lines
5.1 KiB
Python
import cv2
|
|
import numpy as np
|
|
from flask import Flask, Response, request
|
|
import json
|
|
import time
|
|
import sys
|
|
|
|
|
|
max_readings = int(sys.argv[1])
|
|
porta = sys.argv[2]
|
|
url = sys.argv[3]
|
|
arquivoSaida = sys.argv[4]
|
|
|
|
app = Flask(__name__)
|
|
# Função para detecção de objetos usando YOLO e salvar as coordenadas em um arquivo JSON
|
|
def detect_objects(conf_threshold, nms_threshold, _camera_index):
|
|
# Carregar o modelo YOLO e os arquivos de configuração
|
|
model_config = 'C:/Zendion Inc/agrobot_base/Python/models/yolo/yolov3.cfg'
|
|
model_weights = 'C:/Zendion Inc/agrobot_base/Python/models/yolo/yolov3.weights'
|
|
|
|
# Carregar as classes que o modelo pode detectar
|
|
with open('C:/Zendion Inc/agrobot_base/Python/models/yolo/coco.names', 'rt') as f:
|
|
classes = f.read().rstrip('\n').split('\n')
|
|
|
|
net = cv2.dnn.readNetFromDarknet(model_config, model_weights)
|
|
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
|
|
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
|
|
|
|
cap = cv2.VideoCapture(_camera_index, cv2.CAP_DSHOW)
|
|
readings = []
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
# Detecção de Objetos usando YOLO
|
|
blob = cv2.dnn.blobFromImage(frame, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
|
|
net.setInput(blob)
|
|
outs = net.forward(net.getUnconnectedOutLayersNames())
|
|
|
|
# Obter as camadas de saída do modelo YOLO
|
|
output_layers_names = net.getUnconnectedOutLayersNames()
|
|
|
|
# Verificar se há saída válida
|
|
if not output_layers_names:
|
|
print("Nenhuma camada de saída encontrada. Verifique se o modelo foi carregado corretamente.")
|
|
break
|
|
|
|
# Fazer a detecção de objetos no frame
|
|
blob = cv2.dnn.blobFromImage(frame, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
|
|
net.setInput(blob)
|
|
outs = net.forward(output_layers_names)
|
|
|
|
# Mostrar as detecções no frame
|
|
class_ids = []
|
|
confidences = []
|
|
boxes = []
|
|
|
|
for out in outs:
|
|
for detection in out:
|
|
scores = detection[5:]
|
|
class_id = np.argmax(scores)
|
|
confidence = scores[class_id]
|
|
|
|
if confidence > conf_threshold:
|
|
# Obter informações do objeto detectado
|
|
center_x = int(detection[0] * frame.shape[1])
|
|
center_y = int(detection[1] * frame.shape[0])
|
|
w = int(detection[2] * frame.shape[1])
|
|
h = int(detection[3] * frame.shape[0])
|
|
|
|
# Coordenadas do retângulo
|
|
x = int(center_x - w / 2)
|
|
y = int(center_y - h / 2)
|
|
|
|
boxes.append([x, y, w, h])
|
|
confidences.append(float(confidence))
|
|
class_ids.append(class_id)
|
|
|
|
# Aplicar Non-Max Suppression
|
|
indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)
|
|
|
|
# Mostrar as detecções no frame após o NMS
|
|
current_readings = []
|
|
for i in indices:
|
|
box = boxes[i]
|
|
x, y, w, h = box[0], box[1], box[2], box[3]
|
|
class_id = class_ids[i]
|
|
class_name = classes[class_id]
|
|
confidence = confidences[i]
|
|
cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
|
|
cv2.putText(frame, class_name, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
|
|
|
|
|
|
|
|
# Salvar as detecções em um formato desejado
|
|
detection_info = {
|
|
'id': int(i),
|
|
'descricao': class_name,
|
|
'x': int(x),
|
|
'y': int(y),
|
|
'largura': int(w),
|
|
'altura': int(h),
|
|
'confianca': float(confidence)
|
|
}
|
|
#print(detection_info)
|
|
current_readings.append(detection_info)
|
|
|
|
timestamp = time.time()
|
|
readings.append({'timestamp': timestamp, 'objetos': current_readings})
|
|
|
|
# Salvar as detecções em um arquivo JSON
|
|
if len(readings) > max_readings:
|
|
readings.pop(0)
|
|
|
|
if len(readings) == max_readings:
|
|
try:
|
|
with open(arquivoSaida, 'w') as file:
|
|
json.dump({'frames': readings}, file)
|
|
except Exception as e:
|
|
print(f"Erro ao escrever no arquivo: {e}")
|
|
|
|
cv2.imshow('Deteccao de Objetos YOLO', frame)
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
break
|
|
|
|
|
|
|
|
# Convertendo frame para JPEG e enviando via Flask
|
|
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')
|
|
|
|
@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')
|
|
|
|
if __name__ == "__main__":
|
|
app.run(host='0.0.0.0', port=porta)
|