agrobot_base/Python/object-recognize/object-recognize-save-mp-tr...

73 lines
2.4 KiB
Python

import cv2
import json
import time
from flask import Flask, Response
import sys
max_readings = 10
cameraIndex = int(sys.argv[1])
porta = sys.argv[2]
url = sys.argv[3]
arquivoSaida = sys.argv[4]
app = Flask(__name__)
cap = cv2.VideoCapture(cameraIndex)
# Carregando o modelo pré-treinado para detecção de objetos
model_path = 'C:\\Zendion Inc\\agrobot_base\\Python\\models\\coco\\faster_rcnn_inception_v2_coco_2018_01_28\\frozen_inference_graph.pb'
model = cv2.dnn.readNetFromTensorflow(model_path)
readings = []
def generate_frames():
while True:
ret, frame = cap.read()
# Realizando a detecção de objetos no frame
blob = cv2.dnn.blobFromImage(frame, size=(300, 300), swapRB=True, crop=False)
model.setInput(blob)
detections = model.forward()
current_readings = []
for i in range(detections.shape[2]):
confidence = detections[0, 0, i, 2]
if confidence > 0.5: # Defina um limite de confiança adequado
box = detections[0, 0, i, 3:7] * 300
x, y, w, h = box.astype(int)
current_readings.append({'id': i, 'x': x, 'y': y, 'largura': w, 'altura': h})
cv2.rectangle(frame, (x, y), (w, h), (255, 0, 0), 2)
cv2.putText(frame, f'Objeto {i}', (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
timestamp = time.time()
readings.append({'timestamp': timestamp, 'objetos': current_readings})
if len(readings) > max_readings:
readings.pop(0)
if len(readings) == max_readings:
try:
with open('scripts/' + arquivoSaida, 'w') as file:
json.dump({'frames': readings}, file)
except Exception as e:
print(f"Erro ao escrever no arquivo: {e}")
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)
def video_feed():
return Response(generate_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')
if __name__ == "__main__":
# Verificando se os argumentos foram fornecidos corretamente
if len(sys.argv) != 5:
print("Usage: python script.py <cameraIndex> <porta> <url> <arquivoSaida>")
sys.exit(1)
app.run(host='0.0.0.0', port=int(porta))