83 lines
2.8 KiB
Python
83 lines
2.8 KiB
Python
import cv2
|
|
import numpy as np
|
|
|
|
# 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')
|
|
|
|
# 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_CPU) # Use DNN_TARGET_OPENCL para GPU
|
|
|
|
# Configurações adicionais
|
|
conf_threshold = 0.5
|
|
nms_threshold = 0.4
|
|
|
|
# Iniciar a captura de vídeo (altere o índice conforme necessário)
|
|
cap = cv2.VideoCapture(0)
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
# 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
|
|
for i in indices:
|
|
box = boxes[i]
|
|
x, y, w, h = box[0], box[1], box[2], box[3]
|
|
cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
|
|
cv2.putText(frame, classes[class_ids[i]], (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
|
|
|
|
cv2.imshow('Deteccao de Objetos YOLO', frame)
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
break
|
|
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|