Detectando objetos em imagens com o algoritmo YOLO

A detecção de objetos é uma tarefa bastante desafiadora, pois ela envolve não apenas a classificação de um determinado elemento em uma imagem, como também a determinação da sua posição. Diversos algoritmos foram criados ao longo do tempo para lidar com esta tarefa, mas todos sempre têm o problema em comum da lentidão e do esforço computacional necessário para processar cada imagem.

Neste artigo falaremos sobre o YOLO, que usa uma abordagem diferenciada de processamento de imagens que permite ter uma boa confiabilidade e rodar em tempo real.

O algoritmo YOLO (You Only Look Once), que do inglês significa “você olha apenas uma vez”, é um algoritmo de visão computacional muito utilizado para detectar objetos em imagens, ele se destaca devido a sua velocidade, podendo detectar em tempo real, tornando útil para aplicações como carros autônomos, sistemas de segurança, monitoramento de vídeo e muitas outras. Ele funciona basicamente dividindo uma imagem em uma grade de células e a cada célula é atribuída a responsabilidade de prever a presença e as informações dos objetos contidos nela, como as caixas delimitadoras (bounding boxes) e a probabilidade de um objeto pertencer a determinada classe.

Diferente de outros algoritmos de visão computacional que fazem o processamento de uma imagem múltiplas vezes, o YOLO utiliza uma única rede neural convolucional, tratando o problema de detecção como uma tarefa de regressão. As imagens, são compostas de pixels, cada pixel é composto de 3 canais que vão de 0 a 255 cada e representam a intensidade para o vermelho, verde e azul. Ao combinar diferentes intensidades dessas três cores, pode-se obter qualquer cor. Podemos visualizar e manipular esses parâmetros utilizando bibliotecas como o Numpy na linguagem Python.

Uma das operações que podem ser realizadas é a transformação da imagem utilizando um Kernel (Filtro), que consiste em passar outra imagem pequena que contém valores fixos para multiplicar os valores da imagem original. Existem diversos tipos de Kernels e eles são os responsáveis por criar os efeitos em imagens que estão presentes em editores como o desfoque, aprimorar nitidez, realçar bordas, etc. Esse processo é chamado de convolução e o que as redes neurais convolucionais fazem, em resumo, é passar a imagem por vários desses Kernels, projetados especificamente para extrair informações. Então, a partir delas, é possível classificar a imagem.

Configurando o YOLO

Para começar, verifique se o Python está instalado, crie um ambiente virtual e depois baixe as bibliotecas necessárias digitando o comando abaixo no terminal:

Em seguida, crie um arquivo chamado main.py, nele vamos importar a classe YOLO da biblioteca ultralytics e importar o OpenCV:

Agora criaremos uma instância da classe YOLO, passando uma string com o nome yolov8n.pt, que será o nosso modelo. A string passada específica qual é a versão e o tamanho utilizado, no exemplo, o ‘v8’ indica que utilizamos a versão 8 e o ‘n’ o tamanho que é nano. Na primeira vez que for executado, os pesos da rede neural serão baixados:

Outra opção é utilizar o yolov8l.pt, a letra ‘L’ significa que queremos utilizar o modelo grande (‘large’), ele tem uma capacidade maior de dizer a confiança para os objetos, como também prever mais objetos em uma imagem. Com o modelo carregado, basta chamar ele passando o nome da imagem e definindo o parâmetro ‘show’ como true, depois executamos a função waitKey do OpenCV para conseguirmos visualizar a imagem:

Na imagem abaixo você pode ver o resultado do modelo maior detectando pessoas entrando em um ônibus, as mochilas, o skate e até o próprio ônibus foram detectados.

Você pode testar os modelos já treinados com diversas classes em diferentes imagens ou vídeos e analisar o resultado que será gerado. A partir disso, pode projetar sistemas para contagem de pessoas, carros, motos, rastrear objetos, analisar o fluxo de algo e muito mais.

Utilizando a GPU

Antes de começarmos projetos mais complexos, que envolvem o processamento de imagens em vídeos (os frames) ou em tempo real com uma câmera, primeiro temos que configurar para que o algoritmo utilize a placa de vídeo (GPU) ao invés do processador (CPU), pois a GPU possui uma capacidade maior para processar tarefas desse tipo, então dependendo da placa que você possui e do tamanho do modelo que está tentando utilizar, ela vai aumentar significativamente a performance em comparação com o processador.

Para começar a utilizar a GPU, se estiver utilizando uma placa da NVIDIA, tenha certeza que os drivers estão atualizados verificando no site oficial, depois acesse a página do CUDA toolkit da NVIDIA, esse é um conjunto de ferramentas para criar aplicações aceleradas por GPU. O processo para baixar e instalar é simples, basta fazer o download no site, executar o instalador e pressionar ‘next’ até o final. Depois, é necessário fazer o cadastro na plataforma de desenvolvedor da NVIDIA para baixar o cuDNN. Então, basta copiar os arquivos de cada pasta para dentro da pasta com o mesmo nome localizado onde foi instalado o CUDA toolkit. Feito isso, a última etapa é acessar o site do pytorch e baixar a versão da biblioteca que é compatível com a versão do CUDA.

Com essa configuração feita, a partir de agora a GPU deve ser usada no lugar do processador ao rodar o código do YOLO. Se ainda assim não funcionar, verifique se o CUDA toolkit está nas variáveis de ambiente e se não há outra versão do pytorch que seja incompatível atrapalhando, se tiver, desinstale com o código abaixo e instale novamente a versão correta.

Projeto contagem de carros

O projeto que vamos desenvolver agora é simulando uma câmera de segurança que precisa contar carros. Para isso, crie um novo arquivo python e importe a classe YOLO da biblioteca Ultralytics, o OpenCV e a biblioteca Math:

Em seguida, baixe um vídeo de carros passando em uma rua no YouTube e coloque na mesma pasta do projeto com o nome cars.mp4, depois carregue ele com a função VideoCapture do OpenCV, o modelo grande do YOLO e uma fonte do OpenCV

Os modelos já treinados do YOLO, possuem uma certa quantidade de classes que conseguem detectar, já deixei elas prontas em um array abaixo que você pode copiar e colar no seu código:

O vídeo que utilizei foi o da imagem abaixo, note que existem várias estradas por onde os carros podem circular. Supondo que o sistema precisasse detectar os carros que chegam pela via do centro à esquerda, o que precisamos fazer agora é criar uma máscara para que o algoritmo só detecte os carros na estrada que queremos, caso contrário, ele vai detectar todos os carros que estiverem na imagem.

Para criar a máscara, você pode ir em um aplicativo de edição de imagens como o Canva e colar o vídeo propriamente ou um print da tela (que seja do mesmo tamanho do vídeo), depois basta colocar quadrados pretos por cima das partes onde você não quer que o algoritmo ache os carros. Conforme o exemplo abaixo:

Em seguida, você apaga o vídeo que está atrás para ficar só com a máscara e salva como png na pasta do projeto. Note que a máscara cobre parte da estrada onde queremos detectar os carros, deixando apenas a parte que está mais próxima da câmera. O motivo é porque o YOLO tem dificuldade de detectar objetos pequenos, então ele pode detectar com uma baixa confiança ou com a classe errada os carros que estão mais distantes. Ao finalizar, ela deverá ficar da seguinte forma:

No código, você carrega a máscara como uma imagem utilizando a função imread do OpenCV:

Agora, para detectar os carros, precisaremos fazer um loop infinito que vai ficar lendo os frames do vídeo, juntando a máscara e chamando o modelo para detectar os objetos na imagem gerada:

A função cap.read() vai ler os frames, o cv.bitwise_and() vai juntar com a máscara e salvar na variável imgRegion, depois passamos para o modelo e armazenamos o resultado na variável results. Feito isso, ainda dentro do loop, precisamos pegar esses resultados e desenhar na tela a caixa delimitadora, a confiabilidade da previsão e a classe.

Podemos utilizar um loop ‘for’ para iterar sobre cada resultado, cada um deles pode conter múltiplos objetos da classe ‘Boxes’ que são usados pelo Pytorch para armazenar e manipular as bounding boxes, então precisaremos criar outro loop ‘for’ para iterar sobre elas também:

A confiança é obtida a partir da propriedade conf de box e utilizamos a função ceil da biblioteca math para transformar no formato apropriado para visualizarmos. A classe é obtida pela propriedade cls, que vai retornar um número, então passamos esse número para o array com o nome das classes para obter a respectiva string. Utilizando o if podemos selecionar as classes que queremos que o modelo detecte e a confiança mínima. Extraímos as coordenadas do objeto na imagem da propriedade xyxy, transformamos em inteiro e passamos para a função rectangle e putText do opencv para desenhar na tela as bounding boxes, a classe e a confiança.

Agora, fora dos loops ‘for’ mas ainda dentro do loop infinito, temos que desenhar a imagem na tela e chamar a função waitKey para podermos visualizar o algoritmo rodando:

Com o código que escrevemos até aqui, este será o resultado:

Para finalizar o projeto, só o que falta é contar os carros. Para isso, precisamos de um rastreador que guarde a informação que diferencie os carros, por exemplo, atribuindo um id a cada um, dessa forma garantimos que eles não sejam contados mais de uma vez. Para tal, é necessário instalar mais duas bibliotecas, o scikit-image e o filterpy:

Elas serão necessárias para o rastreador, que você pode baixar acessando o GitHub do desenvolvedor abewley no projeto sort. Utilizaremos apenas o arquivo sort.py, que você pode fazer o download e salvar ou copiar o conteúdo e colar dentro de um arquivo com o mesmo nome na pasta do projeto. No começo do arquivo do detector de carros, importe tudo do sort.py e crie um objeto da classe Sort passando o max_age como 20 (que será o número de frames que o algoritmo vai esperar para detectar novamente o mesmo carro se ele sumir):

Agora antes do ‘for’ loop, após a definição da variável results, vamos criar um array numpy vazio, ao qual serão adicionados os resultados da posição e confiança dos resultados. Funciona como se estivéssemos adicionando valores a uma lista normalmente do Python, a diferença é que o numpy chama de stack. Feito isso, basta atualizar os resultados do rastreador, chamando a função update e passando as detecções do array numpy:

Dentro do ‘if’, criamos um array com os parâmetros de posição e confiança e adicionamos os valores para a variável detection utilizando a função vstack do numpy. Quando o loop termina, os valores são atualizados e salvos na variável trackerResults. O que o código que importamos do github vai fazer é gerar um id para cada detecção, que vai se manter ao longo dos frames do vídeo, com isso podemos saber a posição e diferenciar os carros.

Uma forma de fazer essa contagem é desenhar uma linha na tela, fazer um loop pelos resultados do rastreador e para cada resultado você pega os valores de x2 e y2 e diminui de x1 e y1 para achar o tamanho e largura respectivamente. Então, para achar o ponto central da imagem no eixo x e y, diminui a metade da largura de x1 e a metade da altura por y1, respectivamente. Sabendo o ponto central, basta verificar se ele está próximo da linha que foi traçada, se sim, verificar se o número do id  ainda não está presente em um array.

No código, adicionamos antes do loop infinito um array que será onde adicionaremos o id dos carros:

Dentro do loop infinito, depois de trackerResults, traçamos a linha e fazemos o loop para iterar sobre os resultados do rastreador:

Para finalizar, basta mostrar na tela o tamanho do array, que será o número de carros que foram contados:

Conclusão

A inteligência artificial é um tópico que vem crescendo bastante em popularidade nos últimos anos, em especial, a área da visão computacional se destaca pelas inúmeras aplicações práticas que podem ser desenvolvidas como a criação de novos métodos de análise qualitativa de amostras e reações para a indústria farmacêutica ou o desenvolvimento de carros autônomos. Tecnologias essas que há poucos anos poderiam ser consideradas ficção científica, mas que hoje já estão se tornando realidade com o desenvolvimento de novos algoritmos para processamento de imagens em tempo real.

A detecção de objetos em imagens e vídeos é uma área bastante ampla, existem diversos algoritmos para realizar essa tarefa e o YOLO é um deles. Neste artigo exploramos um pouco do potencial dele utilizando os modelos já treinados para criar um sistema que conta carros automaticamente. O YOLO também permite que você crie seu próprio banco de dados e treine para detectar com as classes que você definir, o que é um assunto que poderemos abordar mais detalhadamente em um próximo artigo.