Xây dựng ứng dụng tìm kiếm hình ảnh dùng AI với BridgeTower và Streamlit

Là một Kỹ sư Gen AI, tôi luôn ưu tiên cập nhật các công cụ và công nghệ mới nhất bằng cách liên tục đọc và khám phá những tiến bộ mới. Gần đây, tôi bắt gặp một bài báo nghiên cứu của các nhà nghiên cứu Microsoft giới thiệu BridgeTower, một mô hình giúp đơn giản hóa việc tích hợp khả năng thị giác và ngôn ngữ, giúp các tác vụ đa phương thức dễ triển khai hơn rất nhiều.

Trong hướng dẫn này, tôi sẽ chỉ bạn cách xây dựng một hệ thống truy xuất hình ảnh–văn bản đơn giản nhưng mạnh mẽ bằng cách sử dụng BridgeTower từ thư viện Hugging Face Transformers. Bài hướng dẫn này thân thiện với người mới bắt đầu và sẽ dẫn bạn qua các bước sau:

• Thiết lập các thư viện cần thiết.
• Xử lý hình ảnh và văn bản với BridgeTower.
• Tạo embedding cho cả hình ảnh và văn bản.
• So sánh các embedding này bằng độ tương đồng cosine.

Kết thúc bài hướng dẫn, bạn sẽ có một hệ thống hoàn chỉnh có thể ghép nối hiệu quả hình ảnh với truy vấn văn bản, thể hiện sức mạnh và sự đơn giản của AI đa phương thức.

BridgeTower là gì?

Mô hình BridgeTower, như được giải thích trong bài báo BridgeTower: Building Bridges Between Encoders in Vision-Language Representative Learning của Xiao Xu và cộng sự, nhằm cải thiện cách các mô hình thị giác và ngôn ngữ phối hợp với nhau.

Hãy hình dung một tác vụ nơi mô hình cần hiểu cả hình ảnh và mô tả văn bản, chẳng hạn “một con chó đang chạy trong công viên.” 

Thông thường, mô hình sẽ xử lý hình ảnh và văn bản riêng rẽ bằng các bộ mã hóa thị giác và ngôn ngữ. Mô hình BridgeTower thêm một “cầu nối” giữa các bộ mã hóa này và bộ mã hóa liên phương thức (kết hợp cả thị giác và ngôn ngữ), cho phép chúng tương tác tốt hơn ở mỗi bước.

Nguồn

Ví dụ, trong khi xử lý hình ảnh con chó, bộ mã hóa thị giác có thể phát hiện hình dáng và chuyển động của nó. Đồng thời, bộ mã hóa ngôn ngữ hiểu các từ “đang chạy” và “con chó.” 

Cầu nối cho phép hai phần thông tin này tương tác chi tiết, giúp mô hình kết nối tín hiệu thị giác và văn bản hiệu quả hơn. Nhờ đó, mô hình hoạt động tốt hơn ở các tác vụ như tạo chú thích ảnh hoặc hỏi đáp, mà không cần nhiều tính toán bổ sung hay làm chậm tốc độ.

Cấu trúc thư mục ứng dụng

Chúng ta sẽ giữ mọi thứ gọn gàng và có tổ chức. Thư mục dự án của bạn nên như sau:

/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Được hỗ trợ bởi

• app.py: Xử lý giao diện frontend và tương tác người dùng.
• search.py: Chứa logic tạo embedding và thực hiện tìm kiếm.
• images/: Lưu trữ các tệp hình ảnh đã tải lên.
• requirements.txt: Liệt kê các thư viện cần thiết.

Thiết lập và cài đặt

Trước khi đi vào mã nguồn, hãy thiết lập môi trường để đảm bảo quá trình phát triển và kiểm thử suôn sẻ. Thực hiện các bước sau:

1. Cài đặt phụ thuộc

Đầu tiên, đảm bảo bạn đã cài Python (khuyến nghị phiên bản 3.7 trở lên). Sau đó, chạy lệnh sau trong terminal để cài các thư viện cần thiết:

pip install torch transformers streamlit PillowĐược hỗ trợ bởi

• torch: Khung học sâu dùng để xây dựng và chạy mô hình Bridgetower.
• transformers: Cung cấp các mô hình tiền huấn luyện như Bridgetower để xử lý văn bản và hình ảnh hiệu quả.
• streamlit: Một framework web nhẹ để tạo giao diện người dùng tương tác cho ứng dụng.
• Pillow: Thư viện xử lý ảnh, dùng để xử lý tải lên và thay đổi kích thước hình ảnh.

2. Chuẩn bị ảnh mẫu

Để kiểm thử ứng dụng, hãy đặt một số ảnh mẫu vào thư mục /images. Đảm bảo các ảnh này có định dạng được hỗ trợ như JPG, PNG hoặc JPEG, vì đây là các định dạng phổ biến và tương thích với hầu hết thư viện. 

Hãy dùng đa dạng hình ảnh để đảm bảo mô hình có thể truy xuất kết quả chính xác dựa trên truy vấn văn bản.

Giải thích logic cốt lõi của BridgeTower

Chức năng cốt lõi của ứng dụng được chia thành hai thành phần chính: Backend (search.py) và Frontend (app.py). Mỗi thành phần đóng vai trò quan trọng để ứng dụng vận hành trơn tru.

Logic phía backend (search.py)

Tệp này chứa logic cho embedding văn bản và hình ảnh cùng cơ chế tìm kiếm.

1. Import thư viện

from transformers import BridgeTowerProcessor, BridgeTowerForImageAndTextRetrieval
from PIL import Image
import torch.nn.functional as F
import osĐược hỗ trợ bởi

• BridgeTowerProcessor: Lớp tiện ích chuẩn bị dữ liệu (hình ảnh và văn bản) đúng định dạng cho mô hình BridgeTower. Nó thực hiện các tác vụ như tokenization cho văn bản và tiền xử lý cho hình ảnh.
• BridgeTowerForImageAndTextRetrieval: Lớp mô hình BridgeTower tiền huấn luyện, được thiết kế riêng cho các tác vụ truy xuất hình ảnh–văn bản, nơi bạn tìm ảnh phù hợp nhất dựa trên truy vấn văn bản hoặc ngược lại.
• PIL (Python Imaging Library) dùng để xử lý hình ảnh. Ở đây, nó dùng để mở và thao tác ảnh (ví dụ: đổi kích thước, cắt) trước khi đưa vào mô hình.
• torch.nn.functional cung cấp nhiều hàm như hàm kích hoạt (ReLU, softmax), hàm mất mát (cross-entropy, v.v.) và các phép toán khác trên tensor. Thường được dùng khi làm việc với mạng nơ-ron trong PyTorch.
• os dùng để tương tác với hệ điều hành. Hữu ích khi đọc tệp, kiểm tra thư mục và xử lý đường dẫn (ví dụ: xác định vị trí ảnh hoặc trọng số mô hình).

2. Tải mô hình và 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)Được hỗ trợ bởi

Đoạn này khởi tạo mô hình và processor cần thiết để tạo embedding.

3. Tìm kết quả khớp tốt nhất

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 hỗ trợ bởi

Cách hoạt động:

1. Hàm image_search xử lý truy vấn văn bản để tạo embedding của nó bằng mô hình.
2. Tiến hành xử lý mỗi ảnh để tạo embedding tương ứng của chúng.
3. Các embedding được so sánh bằng logits đầu ra của mô hình.
4. Độ tương đồng cosine được tính giữa embedding văn bản và hình ảnh.
5. Hình ảnh có điểm tương đồng cao nhất được xác định và trả về như kết quả phù hợp nhất.

Giao diện frontend (app.py)

Tệp này dùng Streamlit để cung cấp giao diện thân thiện với người dùng.

1. Import thư viện

import streamlit as st
from search import search_best_match
import osĐược hỗ trợ bởi

2. Xây dựng giao diện

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)Được hỗ trợ bởi

Người dùng có thể tải lên nhiều hình ảnh. Ảnh được lưu trong thư mục images/ và hiển thị lại.

3. Thêm truy vấn tìm kiếm

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)Được hỗ trợ bởi

Người dùng nhập truy vấn tìm kiếm. Ứng dụng sẽ tìm hình ảnh phù hợp nhất và hiển thị nó.

Triển khai đầy đủ: Xây dựng ứng dụng tìm kiếm hình ảnh với BridgeTower và Streamlit

1. search.py - Logic 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Được hỗ trợ bởi

2. app.py - Giao diện 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)Được hỗ trợ bởi

3. Chạy ứng dụng

streamlit run app.pyĐược hỗ trợ bởi

Ví dụ ứng dụng BridgeTower và Streamlit

Đây là ví dụ ứng dụng hoạt động! Một hình ảnh được tải lên và người dùng nhập truy vấn tìm kiếm “cat”. Ứng dụng nhanh chóng phân tích hình ảnh và văn bản, rồi ghép chúng bằng mô hình BridgeTower. 

Ứng dụng đã tìm ra kết quả phù hợp nhất và hiển thị đúng hình ảnh một con mèo. Điều này cho thấy ứng dụng kết nối hình ảnh và văn bản tốt thế nào để đưa ra kết quả chính xác.

Sự tò mò không nên dừng lại ở đây. Tiềm năng của các ứng dụng như thế này vượt xa những gì hiện tại đã triển khai, mở ra nhiều cơ hội thú vị cho tăng trưởng và đổi mới. Khả năng nâng cấp trải nghiệm người dùng, hiệu năng và chức năng là rất lớn. Dưới đây, tôi đề xuất một vài cải tiến có thể đưa các ứng dụng như vậy lên tầm cao mới.

Các cải tiến BridgeTower trong tương lai

Có nhiều cách bạn có thể nâng cấp ứng dụng chúng ta đã tạo, ví dụ: 

Thêm lập chỉ mục FAISS để tăng tốc tìm kiếm hình ảnh

Triển khai lập chỉ mục FAISS để cải thiện tốc độ tìm kiếm hình ảnh trên bộ dữ liệu lớn. Cách này cho phép bạn chuyển ảnh thành vector đặc trưng, xây dựng chỉ mục và thực hiện tìm kiếm tương đồng nhanh chóng, đảm bảo truy xuất ảnh vẫn hiệu quả ngay cả khi dữ liệu tăng trưởng.

Triển khai ứng dụng trên Streamlit Cloud hoặc AWS

Triển khai ứng dụng trên Streamlit Cloud để thiết lập nhanh, dễ dàng, đặc biệt khi bạn làm việc với Python và Streamlit. Hoặc triển khai trên AWS để có nhiều quyền kiểm soát hơn, khả năng mở rộng và xử lý lưu lượng cao, đảm bảo hiệu năng tốt hơn khi ứng dụng mở rộng.

Triển khai cập nhật thời gian thực cho tải lên hình ảnh

Thiết lập cập nhật thời gian thực để khi có ảnh mới được tải lên, chúng xuất hiện ngay cho tất cả người dùng mà không cần tải lại trang. Sử dụng các công nghệ như WebSocket hoặc Socket.IO, bạn có thể đẩy cập nhật đến tất cả người dùng đang kết nối ngay khi ảnh được tải lên, nâng cao tính tương tác và trải nghiệm ứng dụng.

Kết luận

Chúng ta đã xây dựng thành công một ứng dụng tìm kiếm hình ảnh dùng AI với Bridgetower, Streamlit và PyTorch, cho thấy AI có thể đơn giản hóa và nâng cao tương tác với nội dung số như thế nào. 

Dù đã hoạt động đầy đủ, vẫn còn không gian cải tiến, như tối ưu tốc độ tìm kiếm với lập chỉ mục FAISS, triển khai trên các nền tảng như AWS để mở rộng quy mô, và thêm cập nhật thời gian thực để trải nghiệm mượt mà hơn. Dự án này là nền tảng vững chắc cho đổi mới tiếp theo trong các hệ thống tìm kiếm thông minh.

Nếu bạn muốn học thêm, tôi có thể gợi ý các tài nguyên sau: 

• Streamlit Python: Hướng dẫn
• Khóa học Transformer Models with PyTorch
• Cách xây dựng giao diện người dùng cho ứng dụng AI bằng Streamlit và LangChain