Crea una app de búsqueda de imágenes con IA usando BridgeTower y Streamlit

Como ingeniero de GenAI, me esfuerzo por mantenerme al día con las últimas herramientas y tecnologías leyendo y explorando continuamente nuevos avances. Hace poco me topé con un paper de investigadores de Microsoft que presenta BridgeTower, un modelo que simplifica la integración de visión y lenguaje, facilitando mucho la implementación de tareas multimodales.

En esta guía, te mostraré cómo crear un sistema sencillo pero potente de recuperación imagen-texto usando BridgeTower de la biblioteca Hugging Face Transformers. Este tutorial está pensado para principiantes y te guiará por estos pasos:

• Configurar las librerías necesarias.
• Procesar imágenes y texto con BridgeTower.
• Generar embeddings para imágenes y texto.
• Comparar esos embeddings con similitud coseno.

Al final del tutorial, tendrás un sistema plenamente funcional capaz de casar imágenes con consultas de texto de forma eficaz, demostrando la potencia y la sencillez de la IA multimodal.

¿Qué es BridgeTower?

El modelo BridgeTower, tal y como se explica en el paper BridgeTower: Building Bridges Between Encoders in Vision-Language Representative Learning de Xiao Xu y otros autores, busca mejorar cómo colaboran los modelos de visión y lenguaje.

Imagina una tarea en la que el modelo debe entender una imagen y una descripción de texto, como "un perro corriendo por un parque". 

Normalmente, el modelo procesaría imagen y texto por separado usando codificadores de visión y lenguaje. BridgeTower añade un "puente" entre esos codificadores y el codificador cruzado (que combina visión y lenguaje), permitiendo que interactúen mejor en cada paso.

Fuente

Por ejemplo, mientras procesa la imagen del perro, el codificador de visión puede detectar su silueta y movimiento. A la vez, el codificador de lenguaje entiende las palabras "corriendo" y "perro". 

El puente permite que ambas piezas de información interactúen con detalle, ayudando al modelo a conectar mejor las pistas visuales y textuales. Así, el modelo rinde mejor en tareas como subtitulado de imágenes o respuesta a preguntas, sin necesidad de mucha computación extra ni ralentizaciones.

Estructura de carpetas de la app

Mantendremos todo limpio y ordenado. Tu directorio de proyecto debería quedar así:

/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: controla el frontend y las interacciones con la persona usuaria.
• search.py: contiene la lógica para generar embeddings y realizar búsquedas.
• images/: almacena los archivos de imagen subidos.
• requirements.txt: lista las librerías necesarias.

Configuración e instalación

Antes de ponernos con el código, preparemos el entorno para asegurar un desarrollo y pruebas fluidos. Sigue estos pasos:

1. Instala las dependencias

Primero, asegúrate de tener Python instalado (se recomienda la versión 3.7 o superior). Luego, ejecuta este comando en tu terminal para instalar las librerías necesarias:

pip install torch transformers streamlit Pillow

• torch: el framework de deep learning usado para ejecutar el modelo BridgeTower.
• transformers: proporciona modelos preentrenados como BridgeTower para procesar texto e imágenes de forma eficiente.
• streamlit: un framework web ligero para crear una interfaz interactiva para la aplicación.
• Pillow: librería de procesamiento de imágenes para gestionar subidas y redimensionado.

2. Prepara imágenes de muestra

Para probar la aplicación, coloca algunas imágenes de ejemplo en el directorio /images. Asegúrate de que estén en formatos compatibles como JPG, PNG o JPEG, ya que son universales y compatibles con la mayoría de librerías. 

Usa variedad de imágenes para comprobar que el modelo recupera coincidencias en función de la consulta de texto.

Lógica central de BridgeTower explicada

La funcionalidad principal de la aplicación se divide en dos componentes: el backend (search.py) y el frontend (app.py). Cada uno cumple un papel clave para que la app funcione sin fricciones.

Lógica del backend (search.py)

Este archivo contiene la lógica de embeddings de texto e imagen y el mecanismo de búsqueda.

1. Importar librerías

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

• BridgeTowerProcessor: clase de utilidad que prepara los datos (imágenes y texto) en el formato correcto para el modelo BridgeTower. Realiza tareas como la tokenización del texto y el preprocesado de imágenes.
• BridgeTowerForImageAndTextRetrieval: clase de modelo BridgeTower preentrenado diseñada específicamente para tareas de recuperación imagen-texto, donde buscas la imagen más relevante para una consulta de texto o al revés.
• PIL (Python Imaging Library) se usa para manejar imágenes. Aquí sirve para abrir y manipular imágenes (p. ej., redimensionar, recortar) antes de pasarlas al modelo.
• torch.nn.functional ofrece funciones como activaciones (ReLU, softmax), funciones de pérdida (entropía cruzada, etc.) y otras operaciones sobre tensores, habituales al trabajar con redes neuronales en PyTorch.
• os permite interactuar con el sistema operativo: leer archivos, comprobar directorios y gestionar rutas (p. ej., para ubicar imágenes o pesos del modelo).

2. Cargar el modelo y el procesador

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)

Esto inicializa el modelo y el procesador necesarios para generar embeddings.

3. Encontrar la mejor coincidencia

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

Cómo funciona:

1. La función image_search procesa la consulta de texto para generar su embedding con el modelo.
2. Procesa cada imagen para generar sus respectivos embeddings.
3. Los embeddings se comparan usando los logits de salida del modelo.
4. Se calcula la similitud coseno entre los embeddings de texto e imagen.
5. Se identifica y devuelve como mejor coincidencia la imagen con mayor puntuación de similitud.

Interfaz del frontend (app.py)

Este archivo usa Streamlit para ofrecer una interfaz fácil de usar.

1. Importar librerías

import streamlit as st
from search import search_best_match
import os

2. Construir la interfaz

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)

Las personas usuarias pueden subir varias imágenes. Se guardan en la carpeta images/ y se muestran en pantalla.

3. Añadir la consulta de búsqueda

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)

La persona introduce una consulta. La app busca la imagen que mejor coincide y la muestra.

Implementación completa: crea una app de búsqueda de imágenes con BridgeTower y Streamlit

1. search.py - Lógica del 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 - Interfaz en 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. Ejecutar la aplicación

streamlit run app.py

Ejemplo de app con BridgeTower y Streamlit

Aquí tienes un ejemplo de la app en acción. Se subió una imagen y el usuario escribió la consulta de búsqueda "cat". La app analizó rápidamente la imagen y el texto, y los emparejó usando el modelo BridgeTower. 

Encontró la mejor coincidencia y mostró la imagen correcta de un gato. Esto demuestra lo bien que la app conecta imágenes y texto para ofrecer resultados precisos.

La curiosidad no debería quedarse aquí. El potencial de aplicaciones como esta va mucho más allá de lo implementado, abriendo oportunidades muy interesantes de crecimiento e innovación. Las posibilidades para mejorar la experiencia de usuario, el rendimiento y la funcionalidad son enormes. A continuación propongo algunas mejoras que pueden llevar este tipo de apps al siguiente nivel.

Mejoras futuras para BridgeTower

Hay varias formas de mejorar la app que hemos creado, por ejemplo: 

Añadir indexación con FAISS para búsquedas más rápidas

Implementa FAISS para acelerar la búsqueda de imágenes en conjuntos de datos grandes. Podrás convertir imágenes en vectores de características, construir un índice y realizar búsquedas por similitud de forma rápida, manteniendo la eficiencia aunque el dataset crezca.

Desplegar la app en Streamlit Cloud o AWS

Despliega la app en Streamlit Cloud para una puesta en marcha rápida y sencilla, especialmente si trabajas con Python y Streamlit. Como alternativa, súbela a AWS para tener más control, escalabilidad y capacidad de manejar más tráfico, garantizando mejor rendimiento a medida que la app crece.

Implementar actualizaciones en tiempo real para nuevas imágenes

Configura actualizaciones en tiempo real para que, cuando se suba una nueva imagen, aparezca al instante para todas las personas usuarias sin recargar la página. Con tecnologías como WebSockets o Socket.IO podrás enviar actualizaciones a todos los usuarios conectados en cuanto se suba una imagen, mejorando la interactividad y la experiencia.

Conclusión

Hemos construido con éxito una aplicación de búsqueda de imágenes con IA usando BridgeTower, Streamlit y PyTorch, mostrando cómo la IA puede simplificar y mejorar la interacción con contenido digital. 

Aunque es plenamente funcional, hay margen de mejora, como optimizar la velocidad de búsqueda con FAISS, desplegarla en plataformas como AWS para escalar y añadir actualizaciones en tiempo real para una experiencia más fluida. Este proyecto es una base sólida para seguir innovando en sistemas de búsqueda inteligente.

Si te interesa seguir aprendiendo, te recomiendo estos recursos: 

• Streamlit Python: tutorial
• Curso de modelos Transformer con PyTorch
• Cómo crear interfaces de usuario para apps de IA con Streamlit y LangChain: tutorial