Estimativa de pose em tempo real de vídeo em Python com YOLOv7

Introdução

A detecção de objetos é um grande campo em visão computacional e uma das aplicações mais importantes da visão computacional “na natureza”. A partir dele, foi extraída a detecção de pontos-chave (muitas vezes usada para estimativa de pose).

Os pontos-chave podem ser vários pontos – partes de um rosto, membros de um corpo, etc. A estimativa de pose é um caso especial de detecção de pontos-chave – em que os pontos são partes de um corpo humano.

A estimativa de poses é um uso incrível, extremamente divertido e prático da visão computacional. Com ele, podemos acabar com o hardware usado para estimar poses (trajes de captura de movimento), que são caros e difíceis de manejar. Além disso, podemos mapear o movimento de humanos para o movimento de robôs no espaço euclidiano, permitindo movimentos motores de precisão fina sem o uso de controladores, que geralmente não permitem níveis mais altos de precisão. A estimativa de pontos-chave pode ser usada para traduzir nossos movimentos para modelos 3D em AR e VR, e cada vez mais está sendo usada para fazer isso apenas com uma webcam. Finalmente – a estimativa de poses pode nos ajudar em esportes e segurança.

Neste guia, realizaremos uma estimativa de pose em tempo real a partir de um vídeo em Python, usando o modelo YOLOv7 de última geração.

Especificamente, trabalharemos com um vídeo das olimpíadas de inverno de 2018, realizadas em PyeongChang, na Coreia do Sul:

Aljona Savchenko e Bruno Massot fizeram uma performance incrível, incluindo corpos sobrepostos contra a câmera, movimentos rápidos e fluidos e giros no ar. Será uma oportunidade incrível de ver como o modelo lida com situações difíceis de inferir!

YOLO e estimativa de pose

YOLO (Você só olha uma vez) é uma metodologia, bem como uma família de modelos construídos para detecção de objetos. Desde o início em 2015, YOLOv1, YOLOv2 (YOLO9000) e YOLOv3 foram propostos pelo(s) mesmo(s) autor(es) – e a comunidade de aprendizado profundo continuou com avanços de código aberto nos anos contínuos.

Ultralytics' YOLOv5 é um repositório de detecção de objetos de nível industrial, construído sobre o método YOLO. Ele é implementado no PyTorch, ao contrário do C++ para os modelos YOLO anteriores, é totalmente de código aberto e possui uma API lindamente simples e poderosa que permite inferir, treinar e personalizar o projeto de forma flexível. É tão básico que a maioria das novas tentativas de melhorar o método YOLO se baseiam nele.

É assim YOLOR (Você só aprende uma representação) e YOLOv7 que foi construído em cima de YOLOR (mesmo autor) também foram criados!

O YOLOv7 não é apenas uma arquitetura de detecção de objetos - ele fornece novos cabeçotes de modelo, que podem gerar pontos-chave (esqueletos) e realizar segmentação de instâncias, além de apenas regressão de caixa delimitadora, que não era padrão nos modelos YOLO anteriores. Isso não é surpreendente, já que muitas arquiteturas de detecção de objetos foram reaproveitadas para tarefas de segmentação de instância e detecção de ponto-chave anteriormente, devido à arquitetura geral compartilhada, com saídas diferentes dependendo da tarefa.

Mesmo que não seja surpreendente – o suporte à segmentação de instâncias e detecção de pontos-chave provavelmente se tornará o novo padrão para modelos baseados em YOLO, que começaram a superar praticamente todos os outros detectores de dois estágios há alguns anos em termos de precisão e velocidade.

Isso torna a segmentação de instâncias e a detecção de pontos-chave mais rápidas do que nunca, com uma arquitetura mais simples do que os detectores de dois estágios.

O modelo em si foi criado por meio de mudanças arquitetônicas, além de otimizar aspectos do treinamento, apelidado de “bag-of-freebies”, o que aumentou a precisão sem aumentar o custo de inferência.

Instalando o YOLOv7

Vamos começar clonando o repositório para obter o código-fonte:

! git clone https://github.com/WongKinYiu/yolov7.git

Agora, vamos passar para o yolov7 diretório, que contém o projeto, e dê uma olhada no conteúdo:

%cd yolov7
!ls

/content/yolov7
cfg	   figure      output.mp4	 test.py       
data	   hubconf.py  paper		 tools
deploy	   inference   README.md	 train_aux.py
detect.py  LICENSE.md  requirements.txt  train.py
export.py  models      scripts		 utils

Agora, o YOLO deve ser um detector de objetos e não é fornecido com pesos de estimativa de pose por padrão. Queremos baixar os pesos e carregar uma instância de modelo concreto a partir deles. Os pesos estão disponíveis no mesmo repositório do GitHub e também podem ser baixados facilmente pela CLI:

! curl -L https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-w6-pose.pt -o yolov7-w6-pose.pt

 % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0
100  153M  100  153M    0     0  23.4M      0  0:00:06  0:00:06 --:--:-- 32.3M

Uma vez baixado, podemos importar as bibliotecas e métodos auxiliares que usaremos:

import torch
from torchvision import transforms

from utils.datasets import letterbox
from utils.general import non_max_suppression_kpt
from utils.plots import output_to_keypoint, plot_skeleton_kpts

import matplotlib.pyplot as plt
import cv2
import numpy as np

Excelente! Vamos continuar carregando o modelo e criando um script que permite inferir poses de vídeos com YOLOv7 e OpenCV.

Estimativa de pose em tempo real com YOLOv7

Vamos primeiro criar um método para carregar o modelo a partir dos pesos baixados. Vamos verificar qual dispositivo temos disponível (CPU ou GPU):

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

def load_model():
    model = torch.load('yolov7-w6-pose.pt', map_location=device)['model']
    
    model.float().eval()

    if torch.cuda.is_available():
        
        
        model.half().to(device)
    return model

model = load_model()

Dependendo se tivermos uma GPU ou não, ativaremos a meia precisão (usando float16 em vez de float32 nas operações), o que torna a inferência significativamente mais rápida. Observe que é altamente recomendável fazer isso em uma GPU para velocidades em tempo real, pois as CPUs provavelmente não terão o poder de fazê-lo, a menos que sejam executadas em vídeos pequenos.

Vamos escrever um método de conveniência para executar inferência. Aceitaremos imagens como matrizes NumPy (já que é isso que as passaremos mais tarde durante a leitura do vídeo). Primeiro, usando o letterbox() função – vamos redimensionar e preencher o vídeo para uma forma com a qual o modelo possa trabalhar. Isso não precisa ser e não será a forma (resolução) do vídeo resultante!

Em seguida, aplicaremos as transformações, converteremos a imagem para meia precisão (se houver uma GPU disponível), agruparemos e executaremos no modelo:

def run_inference(image):
    
    image = letterbox(image, 960, stride=64, auto=True)[0] 
    
    image = transforms.ToTensor()(image) 
    if torch.cuda.is_available():
      image = image.half().to(device)
    
    image = image.unsqueeze(0) 
    with torch.no_grad():
      output, _ = model(image)
    return output, image

Retornaremos as previsões do modelo, bem como a imagem como tensor. Essas são previsões “rugosas” – elas contêm muitas ativações que se sobrepõem, e nós queremos “limpá-las” usando a Supressão Não Máxima e plotar os esqueletos previstos sobre a própria imagem:

def draw_keypoints(output, image):
  output = non_max_suppression_kpt(output, 
                                     0.25, 
                                     0.65, 
                                     nc=model.yaml['nc'], 
                                     nkpt=model.yaml['nkpt'], 
                                     kpt_label=True)
  with torch.no_grad():
        output = output_to_keypoint(output)
  nimg = image[0].permute(1, 2, 0) * 255
  nimg = nimg.cpu().numpy().astype(np.uint8)
  nimg = cv2.cvtColor(nimg, cv2.COLOR_RGB2BGR)
  for idx in range(output.shape[0]):
      plot_skeleton_kpts(nimg, output[idx, 7:].T, 3)

  return nimg

Com isso em vigor, nosso fluxo geral ficará assim:

img = read_img()
outputs, img = run_inference(img)
keypoint_img = draw_keypoints(output, img)

Para traduzir isso para uma configuração de vídeo em tempo real – usaremos o OpenCV para ler um vídeo e executar esse processo para cada quadro. Em cada quadro, também gravaremos o quadro em um novo arquivo, codificado como um vídeo. Isso necessariamente desacelerará o processo enquanto estamos executando a inferência, exibindo-a e gravando-a – para que você possa acelerar a inferência e a exibição evitando a criação de um novo arquivo e a gravação no loop:

def pose_estimation_video(filename):
    cap = cv2.VideoCapture(filename)
    
    fourcc = cv2.VideoWriter_fourcc(*'MP4V')
    out = cv2.VideoWriter('ice_skating_output.mp4', fourcc, 30.0, (int(cap.get(3)), int(cap.get(4))))
    while cap.isOpened():
        (ret, frame) = cap.read()
        if ret == True:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            output, frame = run_inference(frame)
            frame = draw_keypoints(output, frame)
            frame = cv2.resize(frame, (int(cap.get(3)), int(cap.get(4))))
            out.write(frame)
            cv2.imshow('Pose estimation', frame)
        else:
            break

        if cv2.waitKey(10) & 0xFF == ord('q'):
            break

    cap.release()
    out.release()
    cv2.destroyAllWindows()

A VideoWriter aceita vários parâmetros – o nome do arquivo de saída, o FourCC (quatro códigos de codec, denotando o codec usado para codificar o vídeo), a taxa de quadros e a resolução como uma tupla. Para não adivinhar ou redimensionar o vídeo - usamos a largura e a altura do vídeo original, obtidas através do VideoCapture instância que contém dados sobre o próprio vídeo, como largura, altura, número total de quadros etc.

Agora, podemos chamar o método em qualquer vídeo de entrada:

pose_estimation_video('../ice_skating.mp4')

Isso abrirá uma janela do OpenCV, exibindo a inferência em tempo real. E também gravará um arquivo de vídeo no yolov7 diretório (já que temos cd'd nele):

Conclusão

Neste guia, examinamos o método YOLO, YOLOv7 e a relação entre YOLO e detecção de objetos, estimativa de pose e segmentação de instância. Em seguida, analisamos como você pode instalar e trabalhar facilmente com o YOLOv7 usando a API programática e criar vários métodos de conveniência para facilitar a inferência e a exibição de resultados.

Por fim, abrimos um vídeo usando OpenCV, executamos inferência com YOLOv7 e criamos uma função para realizar estimativa de pose em tempo real, salvando o vídeo resultante em resolução máxima e 30FPS em seu disco local.

Indo além - Aprendizado profundo prático para visão computacional

Sua natureza curiosa faz você querer ir mais longe? Recomendamos verificar nosso Para: “Aprendizado Profundo Prático para Visão Computacional com Python”.

Outro curso de visão computacional?

Não faremos classificação de dígitos MNIST ou moda MNIST. Eles cumpriram sua parte há muito tempo. Muitos recursos de aprendizado estão se concentrando em conjuntos de dados básicos e arquiteturas básicas antes de permitir que arquiteturas de caixa preta avançadas carreguem o fardo do desempenho.

Queremos focar em desmistificação, praticidade, compreensão, intuição e projetos reais. Quero aprender como você pode fazer a diferença? Vamos levá-lo em um passeio desde a maneira como nossos cérebros processam imagens para escrever um classificador de aprendizado profundo de nível de pesquisa para câncer de mama até redes de aprendizado profundo que “alucinam”, ensinando os princípios e a teoria por meio de trabalho prático, equipando você com o know-how e ferramentas para se tornar um especialista na aplicação de aprendizado profundo para resolver a visão computacional.

O que tem dentro?

• Os primeiros princípios da visão e como os computadores podem ser ensinados a “ver”
• Diferentes tarefas e aplicações de visão computacional
• As ferramentas do comércio que facilitarão seu trabalho
• Encontrar, criar e utilizar conjuntos de dados para visão computacional
• A teoria e aplicação das Redes Neurais Convolucionais
• Lidar com mudança de domínio, coocorrência e outros vieses em conjuntos de dados
• Transferir Aprendizagem e utilizar o tempo de treinamento e recursos computacionais de outros para seu benefício
• Construindo e treinando um classificador de câncer de mama de última geração
• Como aplicar uma dose saudável de ceticismo às ideias dominantes e entender as implicações de técnicas amplamente adotadas
• Visualizando o “espaço conceitual” de um ConvNet usando t-SNE e PCA
• Estudos de caso de como as empresas usam técnicas de visão computacional para obter melhores resultados
• Avaliação adequada do modelo, visualização do espaço latente e identificação da atenção do modelo
• Realizando pesquisas de domínio, processando seus próprios conjuntos de dados e estabelecendo testes de modelo
• Arquiteturas de ponta, a progressão das ideias, o que as torna únicas e como implementá-las
• KerasCV – uma biblioteca WIP para criar pipelines e modelos de última geração
• Como analisar e ler artigos e implementá-los você mesmo
• Selecionando modelos dependendo da sua aplicação
• Criando um pipeline de aprendizado de máquina de ponta a ponta
• Paisagem e intuição na detecção de objetos com R-CNNs, RetinaNets, SSDs e YOLO mais rápidos
• Segmentação de instância e semântica
• Reconhecimento de objetos em tempo real com YOLOv5
• Treinamento de detectores de objetos YOLOv5
• Trabalhando com Transformers usando KerasNLP (biblioteca WIP de força da indústria)
• Integrando Transformers com ConvNets para gerar legendas de imagens
• Deepdream
• Otimização de modelo de Deep Learning para visão computacional