Crie um app de busca por imagens com IA usando BridgeTower e Streamlit

Como engenheiro de Gen AI, faço questão de me manter atualizado com as ferramentas e tecnologias mais recentes, lendo e explorando continuamente novas inovações. Recentemente, me deparei com um artigo de pesquisa de pesquisadores da Microsoft apresentando o BridgeTower, um modelo que simplifica a integração de capacidades de visão e linguagem, facilitando muito a implementação de tarefas multimodais.

Neste guia, vou mostrar como criar um sistema simples, mas poderoso, de recuperação imagem-texto usando o BridgeTower da biblioteca Hugging Face Transformers. Este tutorial é amigável para iniciantes e vai conduzir você pelos seguintes passos:

• Configurar as bibliotecas necessárias.
• Processar imagens e texto com o BridgeTower.
• Gerar embeddings para imagens e textos.
• Comparar esses embeddings usando similaridade do cosseno.

Ao final deste tutorial, você terá um sistema totalmente funcional que consegue associar imagens a consultas em texto com eficiência, demonstrando o poder e a simplicidade da IA multimodal.

O que é o BridgeTower?

O modelo BridgeTower, conforme explicado no artigo BridgeTower: Building Bridges Between Encoders in Vision-Language Representative Learning de Xiao Xu e outros autores, busca melhorar como os modelos de visão e linguagem trabalham juntos.

Imagine uma tarefa em que o modelo precisa entender uma imagem e uma descrição em texto, como "um cachorro correndo no parque". 

Normalmente, o modelo processaria a imagem e o texto separadamente usando codificadores de visão e de linguagem. O BridgeTower adiciona uma "ponte" entre esses codificadores e o codificador cruzado (que combina visão e linguagem), permitindo que eles interajam melhor a cada etapa.

Fonte

Por exemplo, ao processar a foto do cachorro, o codificador de visão pode detectar o formato e o movimento do animal. Ao mesmo tempo, o codificador de linguagem entende palavras como "correndo" e "cachorro". 

A ponte permite que essas informações interajam em detalhes, ajudando o modelo a conectar melhor os indícios visuais e textuais. Assim, o desempenho melhora em tarefas como geração de legendas para imagens ou perguntas e respostas, sem exigir muita computação adicional nem tornar o processo mais lento.

Estrutura de pastas do app

Vamos manter tudo limpo e organizado. Seu diretório do projeto deve ficar assim:

/image-search-app/
│
├── app.py          # Frontend (Streamlit UI)
├── search.py       # Backend logic (embeddings & similarity search)
├── images/         # Directory for uploaded images
│   ├── image1.jpg
│   ├── image2.png
│   ├── image3.jpeg
│
└── requirements.txt  # Python dependencies

• app.py: gerencia o frontend e as interações com o usuário.
• search.py: contém a lógica para gerar embeddings e executar as buscas.
• images/: armazena os arquivos de imagem enviados.
• requirements.txt: lista as bibliotecas necessárias.

Configuração e instalação

Antes de mergulhar no código, vamos preparar o ambiente para garantir um processo de desenvolvimento e testes sem atritos. Siga estes passos:

1. Instale as dependências

Primeiro, certifique-se de ter o Python instalado (recomendado a versão 3.7 ou superior). Em seguida, rode o comando abaixo no terminal para instalar as bibliotecas necessárias:

pip install torch transformers streamlit Pillow

• torch: framework de deep learning usado para executar o modelo BridgeTower.
• transformers: fornece modelos pré-treinados como o BridgeTower para processamento eficiente de texto e imagem.
• streamlit: framework web leve para criar uma interface de usuário interativa para o aplicativo.
• Pillow: biblioteca de processamento de imagens, usada para lidar com uploads e redimensionamento.

2. Prepare imagens de teste

Para testar o aplicativo, coloque algumas imagens de exemplo no diretório /images. Certifique-se de que estejam em formatos compatíveis, como JPG, PNG ou JPEG, já que são amplamente reconhecidos e compatíveis com a maioria das bibliotecas. 

Use uma variedade de imagens para garantir que o modelo consiga recuperar correspondências com precisão a partir da consulta em texto.

Lógica central do BridgeTower explicada

A funcionalidade principal do aplicativo está dividida em dois componentes: o backend (search.py) e o frontend (app.py). Cada um tem um papel essencial para que tudo funcione sem fricção.

Lógica do backend (search.py)

Este arquivo contém a lógica de embeddings de texto e imagem e o mecanismo de busca.

1. Importe as bibliotecas

from transformers import BridgeTowerProcessor, BridgeTowerForImageAndTextRetrieval
from PIL import Image
import torch.nn.functional as F
import os

• BridgeTowerProcessor: classe utilitária que prepara os dados (imagens e texto) no formato correto para o modelo BridgeTower. Faz tarefas como tokenização de texto e pré-processamento de imagens.
• BridgeTowerForImageAndTextRetrieval: classe de modelo BridgeTower pré-treinada, projetada especificamente para tarefas de recuperação imagem-texto, em que você busca a imagem mais relevante com base em uma consulta textual (ou o inverso).
• PIL (Python Imaging Library) é usada para manipular imagens. Aqui, serve para abrir e ajustar imagens (por exemplo, redimensionar, cortar) antes de enviá-las ao modelo.
• torch.nn.functional fornece diversas funções, como ativações (ReLU, softmax), funções de perda (cross-entropy etc.) e outras operações com tensores. É muito usada ao trabalhar com redes neurais no PyTorch.
• os é usado para interagir com o sistema operacional. Ajuda a ler arquivos, checar diretórios e lidar com caminhos (por exemplo, para localizar imagens ou pesos de modelos).

2. Carregue o modelo e o processor

processor = BridgeTowerProcessor.from_pretrained("BridgeTower/bridgetower-base-itm-mlm")
model = BridgeTowerForImageAndTextRetrieval.from_pretrained("BridgeTower/bridgetower-base-itm-mlm")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

Isso inicializa o modelo e o processor necessários para gerar embeddings.

3. Encontre a melhor correspondência

def image_search(query, image_dir):
   image_files = [os.path.join(image_dir, img) for img in os.listdir(image_dir) if img.lower().endswith(('jpg', 'jpeg', 'png'))]
   scores = {}

   for img_path in image_files:
       image = Image.open(img_path).convert("RGB")
       encoding = processor(text=query, images=image, return_tensors="pt").to(device)
       outputs = model(**encoding)
       scores[img_path] = outputs.logits[0, 1].item()

   best_match = max(scores, key=scores.get)
   return best_match

Como funciona:

1. A função image_search processa a consulta em texto para gerar seu embedding com o modelo.
2. Ela processa cada imagem para gerar seus respectivos embeddings.
3. Os embeddings são comparados usando os logits de saída do modelo.
4. A similaridade do cosseno é calculada entre os embeddings de texto e imagem.
5. A imagem com maior pontuação de similaridade é identificada e retornada como a melhor correspondência.

Interface do frontend (app.py)

Este arquivo usa o Streamlit para fornecer uma interface amigável.

1. Importe as bibliotecas

import streamlit as st
from search import search_best_match
import os

2. Construa a interface

st.title("🖼️ AI-Powered Image Search App")
uploaded_files = st.file_uploader("Upload images", accept_multiple_files=True, type=['jpg', 'jpeg', 'png'])

image_folder = 'images'
os.makedirs(image_folder, exist_ok=True)

if uploaded_files: 
    for uploaded_file in uploaded_files:
        with open(os.path.join(image_folder, uploaded_file.name), 'wb') as f:
            f.write(uploaded_file.getbuffer())
        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)

Os usuários podem enviar várias imagens. As imagens são salvas na pasta images/ e exibidas.

3. Adicione a consulta de busca

query = st.text_input("Enter your search query:")

if query and uploaded_files:
    best_match, score = search_best_match(image_folder, query)
    st.write("### Most Relevant Image:")
    st.image(best_match, caption=f"Score: {score:.4f}", use_column_width=True)

Os usuários digitam a consulta de busca. O app encontra a imagem mais relevante e a exibe.

Implementação completa: crie um app de busca por imagens com BridgeTower e Streamlit

1. search.py - Lógica do backend

from transformers import BridgeTowerProcessor, BridgeTowerForImageAndTextRetrieval
from PIL import Image
import torch
import torch.nn.functional as F
import os

# Load Bridgetower processor and model
processor = BridgeTowerProcessor.from_pretrained("BridgeTower/bridgetower-base-itm-mlm")
model = BridgeTowerForImageAndTextRetrieval.from_pretrained("BridgeTower/bridgetower-base-itm-mlm")

# Search for the best matching image
def image_search(query, image_dir):
   image_files = [os.path.join(image_dir, img) for img in os.listdir(image_dir) if img.lower().endswith(('jpg', 'jpeg', 'png'))]
   scores = {}

   for img_path in image_files:
       image = Image.open(img_path).convert("RGB")
       encoding = processor(text=query, images=image, return_tensors="pt").to(device)
       outputs = model(**encoding)
       scores[img_path] = outputs.logits[0, 1].item()

   best_match = max(scores, key=scores.get)
   return best_match

2. app.py - Interface em Streamlit

import streamlit as st
from search import image_search
import os

# Title of the app
st.title("🖼️ Image Search Application with BridgeTower")

# Image upload section
uploaded_files = st.file_uploader("Upload your images", accept_multiple_files=True, type=['jpg', 'jpeg', 'png'])

# Create image folder if not exists
image_folder = 'images'
os.makedirs(image_folder, exist_ok=True)

if uploaded_files:
    for uploaded_file in uploaded_files:
        with open(os.path.join(image_folder, uploaded_file.name), 'wb') as f:
            f.write(uploaded_file.getbuffer())
        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)

# Text input for query
query = st.text_input("Enter your search query:")

# Display search results
if query and uploaded_files:
    best_match, score = image_search(image_folder, query)
    st.write("### Most Relevant Image:")
    st.image(best_match, caption=f"Similarity Score: {score:.4f}", use_column_width=True)

3. Rode o aplicativo

streamlit run app.py

Exemplo do app com BridgeTower e Streamlit

Aqui vai um exemplo do app em ação! Uma imagem foi enviada e o usuário digitou a consulta de busca “cat”. O app analisou rapidamente a imagem e o texto e os relacionou usando o modelo BridgeTower. 

Ele encontrou a melhor correspondência e exibiu a imagem correta de um gato. Isso mostra como o app conecta imagens e texto para entregar resultados precisos.

A curiosidade não deve parar por aqui. O potencial de aplicativos como este vai muito além do que está implementado hoje, abrindo oportunidades empolgantes de crescimento e inovação. As possibilidades de aprimorar a experiência do usuário, o desempenho e as funcionalidades são enormes. Abaixo, proponho alguns incrementos que podem levar soluções como esta a outro nível.

Melhorias futuras para o BridgeTower

Há várias formas de melhorar o app que criamos, por exemplo: 

Adicionar indexação FAISS para buscas mais rápidas

Implemente indexação com FAISS para acelerar a busca de imagens em grandes conjuntos de dados. Assim, você converte imagens em vetores de características, constrói um índice e executa buscas por similaridade rapidamente, mantendo a eficiência mesmo com o crescimento do dataset.

Fazer deploy do app no Streamlit Cloud ou na AWS

Faça o deploy no Streamlit Cloud para uma configuração rápida e simples, especialmente se você trabalha com Python e Streamlit. Alternativamente, use a AWS para ter mais controle, escalabilidade e capacidade de lidar com mais tráfego, garantindo melhor performance conforme o app cresce.

Implementar atualizações em tempo real para uploads

Configure atualizações em tempo real para que, quando novas imagens forem enviadas, elas apareçam instantaneamente para todos os usuários, sem precisar recarregar a página. Usando tecnologias como WebSockets ou Socket.IO, você pode enviar atualizações a todos os usuários conectados assim que uma imagem é adicionada, elevando a interatividade e a experiência do app.

Conclusão

Construímos com sucesso um aplicativo de busca por imagens com IA usando BridgeTower, Streamlit e PyTorch, mostrando como a IA pode simplificar e melhorar a interação com conteúdo digital. 

Embora totalmente funcional, ainda há espaço para evoluir: otimizar a velocidade de busca com FAISS, fazer deploy em plataformas como AWS para escalar e adicionar atualizações em tempo real para uma experiência mais fluida. Este projeto é uma base sólida para inovar ainda mais em sistemas inteligentes de busca.

Se você quiser se aprofundar, recomendo estes recursos: 

• Streamlit Python: tutorial
• Curso: Transformer Models with PyTorch
• Como criar interfaces para apps de IA usando Streamlit e LangChain