Estimativa de Velocidade de Carros com YOLO

Nesta postagem você entenderá o processo de estimativa de velocidade de objetos em um vídeo a partir de sua detecção, rastreamento e medição através de um limite definido.

Trabalhando com YOLO

Para utilizar a YOLO, necessitamos primeiramente gerar um modelo para que ele possa ser carregado e, posteriormente, utilizado.

Para isso, quatro etapas são necessárias: o Setup, a Validação, o Treinamento e a Exportação.

Com os comandos abaixo, podemos realizar a primeira etapa do processo:

# Instalação de pacote
%pip install ultralytics

# Importação de biblioteca necessária
import ultralytics
ultralytics.checks()

O processo de validação é realizado com o conjunto de dados COCO. Primeiramente, realizaremos o download e a descompactação do dataset através dos seguintes comandos:

# Importação de biblioteca necessária
import torch

# Download dataset
torch.hub.download_url_to_file('https://ultralytics.com/assets/coco2017val.zip', 'tmp.zip')

# Unzip
!unzip -q tmp.zip -d datasets && rm tmp.zip  

Agora realizaremos a validação de fato. Após a execução do comando, o nosso modelo já é salvo automaticamente.

# Validação YOLOv8n no COCO128 val
!yolo val model=yolov8n.pt data=coco128.yaml

Treinaremos o modelo para que ele realize as detecções de maneira mais assertiva possível:

# Treinamento YOLOv8n no COCO128 para 3 epochs
!yolo train model=yolov8n.pt data=coco128.yaml epochs=3 imgsz=640

Finalmente, exportaremos nosso modelo validado e treinado:

# Exportação de modelo
!yolo export model=yolov8n.pt format=torchscript

Desenvolvendo a aplicação

Iniciaremos nossa aplicação com a importação das bibliotecas necessárias. A YOLO possui um módulo dedicado à estimativa de velocidade de objetos, então nada mais lógico do que utilizá-lo. Também utilizaremos o OpenCV, que nos auxiliará na manipulação dos vídeos e imagens.

# Importação de bibliotecas necessárias
from ultralytics import YOLO
from ultralytics.solutions import speed_estimation
import cv2

Já no início, importaremos nosso modelo treinado e validado da YOLO, bem como salvar em uma variável uma lista com os nomes das classes que a YOLO é capaz de detectar, como, por exemplo, carro (car), pessoa (person), etc.

# Importação de modelo e classes
model = YOLO("yolov8n.pt")
names = model.model.names

Com o OpenCV, iremos iniciar a captura do vídeo e monitorar se algum erro ocorre nesse processo. Também realizaremos a leitura de algumas propriedades, como FPS, largura e altura. Também configuraremos uma variável para realizar o salvamento do vídeo processado, em MP4, com as informações capturadas.

# Variável para captura do vídeo
cap = cv2.VideoCapture("video.mp4")

# Verificação de erro na leitura
assert cap.isOpened()

# Captura de propriedades do vídeo
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))

# Configuração de salvamento do vídeo
video_writer = cv2.VideoWriter("speed_estimation.avi",
                               cv2.VideoWriter_fourcc(*'mp4v'),
                               fps,
                               (w, h))

Também definiremos em uma variável os pontos de nossa linha de medição, para que, quando o objeto ultrapasse esse segmento de reta, seja calculado sua velocidade.

# Definição do segmento de reta (valores dependem do vídeo utilizado)
line_pts = [(1100, 470), (1700, 600)]

Neste momento, instanciaremos o módulo de estimativa de velocidade de objetos da YOLO, passando como argumentos o segmento de reta criado, os nomes das classes e a visualização em tempo real.

# Instanciamento de módulo 
speed_obj = speed_estimation.SpeedEstimator()
speed_obj.set_args(reg_pts=line_pts,
                   names=names,
                   view_img=True)

Também criaremos uma lista, com o primeiro valor sendo zero. Este componente será importante para registrarmos as velocidades dos carros para exibição em tela.

# Criação de lista
lista = [0]

Iniciaremos o loop de leitura do vídeo, lendo os frames e realizando a detecção e rastreamento dos objetos. Logo após, através do módulo de estimativa de velocidade, mediremos a velocidade em si. Uma observação importante é que tivemos que realizar uma pequena alteração no módulo SpeedEstimator() para que ele retornasse, além das informações padrão, o valor numérico da velocidade para exibição.

# Início do loop de leitura dos vídeos
while cap.isOpened():
	
    # Leitura dos frames
    ret, frame = cap.read()

    # Verificação da leitura dos frames
    if not ret:
        print("Video frame is empty or video processing has been successfully completed.")
        break
    # Detecção e rastreamento dos objetos
    tracks = model.track(frame, persist=True, show=False)

    # Estimativa de velocidade
    frame, velocidade = speed_obj.estimate_speed(frame, tracks)

Para finalizar, realizaremos uma pequena lógica para que os últimos valores medidos sempre aparecem na tela.

    # Lógica para exibição da velocidade
    if velocidade != 0:
        lista.append(velocidade) 
        # Desenhar a velocidade no quadro
        cv2.putText(frame, "SPEED: {}".format(lista[-1]), (10, 60),
            cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 2)
    
    else:
        cv2.putText(frame, "SPEED: {}".format(lista[-1]), (10, 60),
                cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 2)
    
    # Gravar vídeo processado
    video_writer.write(frame)

# Finalizações
cap.release()
video_writer.release()
cv2.destroyAllWindows()

Conclusão

É importante lembrar que essa estimativa de velocidade de objetos deve ser implementada com uma calibração de câmera devido ao fato que o cálculo é realizado através das distâncias euclidianas da imagem, ou seja, não se leva em consideração angulações e inclinações, apenas a movimentação dos pontos em um plano 2D. Mesmo assim, conseguimos desenvolver uma aplicação útil que pode ser utilizada em diversos setores para monitoramento de velocidades e medições.