Neste último post da série de visão computacional iremos construir um detector de objetos, em tempo real, utilizando nossos queridinhos: Python, OpenCV e, dessa vez, a YOLO pré-treinada no COCO Dataset.
You Only Look Once (YOLO) + Common Objects in Context (COCO) Dataset
A YOLO é o estado da arte dos sistemas de detecção de objetos em tempo real, extremamente rápida e precisa,
O COCO dataset é um conjunto de dados de ampla escala que contém a classificação de 80 tipos de objetos diferentes com mais de 330K imagens para realização da segmentação de objetos.
Combinados, formam uma poderosa ferramenta para detecção e classificação de objetos em tempo real. Deixo alguns links a seguir se desejar conhecer mais sobre a YOLO e o COCO Dataset.
O Projeto
Nesse projeto iremos utilizar algumas técnicas apresentadas nos artigos anteriores dessa série, por isso não entraremos em detalhes sobre elas e vamos nos concentrar na parte da detecção e identificação dos objetos.
Dessa vez iremos implementar CLI (Command Line Interface), em que o usuário pode informar, através do terminal, a câmera que será utilizada ou um arquivo de vídeo para ser analisado. As funcionalidades que iremos implementar são:
- Inicializar o streaming de vídeo - ou enviar um arquivo de vídeo para ser analisado - à escolha do usuário;
- Realizar pré-processamento da imagem para reduzir os ruídos e;
- Detectar os objetos que desejamos - no nosso caso optamos por detectar somente: pessoa(0), monitor(62), mouse(64), controle remoto(65) e teclado(66).
Preparando nossa aplicação
Primeiramente importaremos as bibliotecas necessárias e configuraremos as constantes necessárias para o correto funcionamento.
import numpy as np
import argparse
import os
from os.path import join, dirname
import cv2
import time
from imutils.video import VideoStream, FileVideoStream
from imutils.video import FPS
Além das bibliotecas que utilizaremos para a captura e processamento das imagens cv2
, imutils
e numpy
, importamos também a biblioteca os
para ler o sistema arquivos.
Como iremos construir um CLI, importamos também a biblioteca argparse
permitindo que o usuário informe os parâmetros necessários via terminal.
Em seguida configuramos algumas constantes:
- Nível de confiança mínimo
- Non-maximum Suppression threshold - um limite para filtrar as "caixas" que destacarão os objetos detectados, evitando múltiplas detecções para um mesmo objeto.
- O diretório que contem nosso modelo treinado e os nomes das classes.
CONFIDENCE_MIN = 0.4
NMS_THRESHOLD = 0.2
MODEL_BASE_PATH = join(dirname(__file__), "yolo-coco")
Em seguida iremos definir algumas funções:
build_parser
: essa função irá construir nosso parser, que possibilitará que nossa aplicação receba argumentos através do terminal:
def build_parser():
parser = argparse.ArgumentParser()
parser.add_argument("-i", "--input", required=True, help="Endereço do streaming, camera index ou caminho do arquivo")
return parser
2. load_classes
: essa função irá extrair os nomes das classes a partir do arquivo coco.names
- retornará os labels e uma cor aleatória para cada label existente no arquivo:
def load_classes():
with open(os.path.sep.join([MODEL_BASE_PATH, 'coco.names'])) as f:
labels = f.read().strip().split('\n')
# gerar cores para cada label
np.random.seed(42)
colors = np.random.randint(0, 255, size=(len(labels), 3), dtype='uint8')
return labels, colors
3. load_model
: irá carregar o modelo treinado no cv2
:
def load_model():
net = cv2.dnn.readNetFromDarknet(
os.path.sep.join([MODEL_BASE_PATH, 'yolov3.cfg']),
os.path.sep.join([MODEL_BASE_PATH, 'yolov3.weights']))
return net
4. extract_layers
: essa função irá extrair as camadas não conectadas da arquitetura YOLO
def extract_layers():
ln = net.getLayerNames()
ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
return ln
5. start_streaming
: responsável por instanciar e inicializar o streaming (ou carregar o arquivo) de vídeo:
def start_streaming(streaming_path):
if os.path.isfile(streaming_path):
vs = FileVideoStream(streaming_path)
elif streaming_path.isnumeric:
vs = VideoStream(int(streaming_path))
else:
vs = VideoStream(streaming_path)
vs.start()
return vs
Agora que nossas funções estão definidas vamos para o código principal, irei separar os blocos que dizem respeito à mesma funcionalidade e comentarei brevemente sobre cada parte no decorrer do artigo:
if __name__ == '__main__':
parser = build_parser()
streaming_path = vars(parser.parse_args())['input']
print("[+] Carregando labels das classes treinadas...")
labels, colors = load_classes()
print("[+] Carregando o modelo YOLO treinado")
net = load_model()
ln = extract_layers()
print("[+] Iniciando a recepção do streaming...")
vs = start_streaming(streaming_path)
time.sleep(1)
fps = FPS().start()
Nessa primeira parte estamos fazendo o setup do nosso CLI:
- construir o parser,
- capturar a entrada do usuário,
- carregar as classes e cores,
- carregar o modelo,
- extrair as camadas desconectadas da arquitetura YOLO e;
- inicializar o streaming de vídeo.
while True:
frame = vs.read()
# redimensionar os frames
# frame = cv2.resize(frame, None, fx=0.2, fy=0.2)
# capturar a largura e altura do frame
(H, W) = frame.shape[:2]
# construir um container blob e fazer uma passagem na YOLO
blob = cv2.dnn.blobFromImage(frame, 1 / 255.0, (W, H), swapRB=True, crop=False)
net.setInput(blob)
layer_outputs = net.forward(ln)
# criar listas com boxes, nível de confiança e ids das classes
boxes = []
confidences = []
class_ids = []
Nosso loop infinito que irá capturar o frame do stream, construir o blob a partir do frame e passar pela YOLO. Iremos também criar as listas de boxes, níveis de confiança e ids das classes, mais adiante utilizaremos essas listas para exibir a localização, nível de confiança e classe dos objetos detectados.
for output in layer_outputs:
for detection in output:
scores = detection[5:]
class_id = np.argmax(scores)
confidence = scores[class_id]
# filtar pelo threshold da confiança
# selecionar somente se for pessoa, monitor, teclado, mouse ou controle remoto
# verificar arquivo coco.name caso precise utilizar outras classes
if confidence > CONFIDENCE_MIN and class_id in [0, 62, 64, 65, 66]:
box = detection[0:4] * np.array([W, H, W, H])
(center_x, center_y, width, height) = box.astype("int")
x = int(center_x - (width / 2))
y = int(center_y - (height / 2))
boxes.append([x, y, int(width), int(height)])
confidences.append(float(confidence))
class_ids.append(class_id)
Nesse trecho estamos percorrendo as saídas e detecções realizadas pela YOLO e filtrando-as tanto pelo nível de confiança mínimo (que definimos na etapa de preparação) quanto pelas classe que desejamos detectar.
Para as detecções que se enquadram nesses critérios, iremos adicionar as respectivas informações nas listas que criamos anteriormente:
-
boxes
- coordenadas da caixa que destaca cada objeto detectado; confidence
- os níveis de confiança de cada objetoclass_ids
- as classes de cada um dos objetos detectados
# eliminar ruido e redundâncias aplicando non-maxima suppression
new_ids = cv2.dnn.NMSBoxes(boxes, confidences, CONFIDENCE_MIN, NMS_THRESHOLD)
if len(new_ids) > 0:
for i in new_ids.flatten():
(x, y) = (boxes[i][0], boxes[i][1])
(w, h) = (boxes[i][2], boxes[i][3])
# plotar retângulo e texto das classes detectadas no frame atual
color_picked = [int(c) for c in colors[class_ids[i]]]
cv2.rectangle(frame, (x, y), (x + w, y + h), color_picked, 2)
text = "{}: {:.4f}".format(labels[class_ids[i]], confidences[i])
cv2.putText(frame, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color_picked, 2)
Como a YOLO não aplica o Non-maximum Supression (NMS) para nós, iremos aplicá-lo para remover a sobreposição das "caixas" que detectam os objetos, mantendo apenas a detecção mais confiável.
Para cada uma dessas detecções (já com o filtro NMS aplicado) iremos extrair os pontos das caixas e plotar o retângulo - ao redor do objeto com o texto e o nível de confiança - no frame atual.
Em seguida vamos exibir o frame atualizado e aguardar o usuário pressionar a tecla de saída (esc):
# exibir o frame atual
cv2.imshow('Frame', frame)
# sair caso seja pressionada a tecla ESC
c = cv2.waitKey(1)
if c == 27:
break
# atualiza o fps
fps.update()
Por fim, eliminamos os processos e janelas:
# eliminar processos e janelas
fps.stop()
cv2.destroyAllWindows()
vs.stop()
Demonstração
Conclusão
Apesar de ser um projeto simples, como conseguimos extrair diversos objetos diferentes podemos imaginar diversas aplicações, tais como: em carros autônomos, detecção de invasão em espaços restritos, identificar utilização de objetos proibido em determinado local, entre outras.