Membuat Aplikasi Pencarian Gambar Dengan Python Menggunakan CLIP Dan Streamlit

Aplikasi seperti Photo Google memungkinkan untuk mencari gambar di handphone Anda menggunakan kueri text. Aspek yang luar biasa ialah aplikasi tidak mengharuskan Anda memberi label pada gambar berdasarkan kontennya. Misalkan, Anda bisa mencari kucing atau sup di aplikasi Photo Google Anda dan mendapatkan hasil yang relevan meskipun gambar Anda tidak memiliki deskripsi text.

Bagaimana aplikasi bisa melakukan ini? Aplikasi seperti ini memahami keterkaitan di antara deskripsi semantik sesuatu episode dan content gambar dari episode tersebut. Di blog ini, saya akan mendemonstrasikan bagaimana Anda dapat menulis aplikasi Pencarian Gambar Anda sendiri dengan Python. Kami akan menggunakan model Pembelajaran Mesin pra-pelatihan yang disebut CLIP yang telah memahami representasi text/gambar bersama yang kami butuhkan. Kami akan menggunakan Streamlit untuk menyajikan aplikasi.

CLIP

Contrastive Language-Image Pretraining (CLIP) ialah model multimodal text/gambar yang populer berdasarkan makalah oleh Radford et al (2021). Model CLIP dilatih pada 400 Juta pasang contoh teks-gambar yang didapat dari internet. Untuk aplikasi kami, kami akan menggunakan model pra-pelatihan untuk mencocokkan istilah pencarian text kami dengan database gambar.

Streamlit

Streamlit ialah framework Python populer yang ditujukan untuk pengembangan aplikasi Machine Learning. Streamlit sebagian besar menangani elemen design estetika pengembangan aplikasi yang memungkinkan kami untuk konsentrasi khususnya pada aspek Pembelajaran Mesin.

Pengembangan Aplikasi

Aplikasi ini terdiri dari 2 skrip:

get_embeddings.py : Dalam skrip ini, kami menyandikan gambar ke embedding menggunakan encoder gambar model CLIP. Penyematan ialah representasi vector dari input yang menyandikan content preskriptifnya.
app.py : Ini ialah program streamlit yang mengimplementasikan fungsi pencarian gambar. Penyematan text diperoleh untuk istilah pencarian input dibandingkan output penyisipan gambar dari langkah pertama. Hasil yang paling sama selanjutnya disajikan dalam format grid.

Code untuk skrip get_embeddings.py dimasukkan di bawah ini. Anda perlu menginstal CLIP menggunakan petunjuk di https://github.com/openai/CLIP.

import os
import clip
import torch
from torch.utils.data import Dataset, DataLoader
import PIL
import pickle
from tqdm import tqdm

class Images(Dataset):
        """Images dataset"""
        
        def __init__(self, image_list, transform):
            """
            Args:
                image_list: List of image paths.
                transform : Transform to be applied on a sample.
            """
            self.image_list = image_list
            self.transform = transform
        
        def __len__(self):
            return len(self.image_list)
        
        def __getitem__(self, idx):
            image_path = self.image_list[idx]
            image = PIL.Image.open(image_path)
            image = self.transform(image)
            data = {'image':image, 
                    'img_path': image_path}
            return data
        
if __name__ == '__main__':
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model, preprocess = clip.load('ViT-B/32', device, jit=False)
    print(f'Device used: {device}')
    
    folder_path = '<Enter folder location with your images here>'
    image_list = [folder_path + file for file in os.listdir(folder_path)]
        
    print('Attempting to open images...')
    cleaned_image_list = []
    for image_path in image_list:
        try:
            PIL.Image.open(image_path)
            cleaned_image_list.append(image_path)
        except:
            print(f"Failed for {image_path}")
    
    print(f"There are {len(cleaned_image_list)} images that can be processed")    
    dataset = Images(cleaned_image_list,preprocess)
    
    dataloader = DataLoader(dataset, 
                            batch_size=256,
                            shuffle=True)
    
    print("Processing images...")
    image_paths = []
    embeddings = []
    for data in tqdm(dataloader):
        with torch.no_grad():
            X = data['image'].to(device)
            image_embedding = model.encode_image(X)
            img_path = data['img_path']
            image_paths.extend(img_path)
            embeddings.extend([torch.Tensor(x).unsqueeze(0).cpu() for x in image_embedding.tolist()])

    image_embeddings = dict(zip(image_paths,embeddings))
    
    # save to pickle file for the app
    print("Saving image embeddings")
    with open('embeddings.pkl','wb') as f:
        pickle.dump(image_embeddings,f)

Kelas Gambar mewariskan dari kelas Dataset pytorch dan memberikan petunjuk tentang cara mengonversi track gambar menjadi tensor pytorch. Dataset digunakan dalam fungsi utama untuk membuat Dataloader pytorch yang memungkinkan kumpulan gambar untuk diproses bersama-sama. Karakter vector dari operasi ini membuatnya bisa lebih cepat daripada memproses gambar satu per satu.

Code memfilter semua track di bawah directory yang ditetapkan dalam variabel folder_path untuk memastikan gambar bisa dibaca oleh PIL. Ini membantu dalam menghapus file-file yang berantakan seperti file.html yang merupakan artefak pengunduhan umum. Setelah embedding dibuat, embedding disimpan dalam file pickles untuk absorbed oleh aplikasi.

Code untuk skrip app.py diberikan di bawah ini.

import streamlit as st
import pandas as pd
import clip
import torch
from sklearn.metrics.pairwise import cosine_similarity
import pickle

device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load('ViT-B/32', device)

# load embeddings from file
with open('embeddings.pkl','rb') as f:
    image_embeddings = pickle.load(f)

st.header('Image Search App')
search_term = 'a picture of ' + st.text_input('Search: ')
search_embedding = model.encode_text(clip.tokenize(search_term).to(device)).cpu().detach().numpy()

st.sidebar.header('App Settings')
top_number = st.sidebar.slider('Number of Search Results', min_value=1, max_value=30)
picture_width = st.sidebar.slider('Picture Width', min_value=100, max_value=500)

df_rank = pd.DataFrame(columns=['image_path','sim_score'])

for path,embedding in image_embeddings.items():
    sim = cosine_similarity(embedding,
                            search_embedding).flatten().item()
    df_rank = pd.concat([df_rank,pd.DataFrame(data=[[path,sim]],columns=['image_path','sim_score'])])
df_rank.reset_index(inplace=True,drop=True)

df_rank.sort_values(by='sim_score',
                    ascending=False,
                    inplace=True,
                    ignore_index=True)

# display code: 3 column view
col1, col2, col3 = st.columns(3)

df_result = df_rank.head(top_number)
for i in range(top_number):
    if i % 3 == 0:
        with col1:
            st.image(df_result.loc[i,'image_path'],width=picture_width)
    elif i % 3 == 1:
        with col2:
            st.image(df_result.loc[i,'image_path'],width=picture_width)
    elif i % 3 == 2:
        with col3:
            st.image(df_result.loc[i,'image_path'],width=picture_width)

Skrip aplikasi memuat penyematan gambar yang disimpan dari cara sebelumnya. Itu mengambil istilah pencarian yang dimasukkan pengguna dari bilah pencarian dan menggunakannya untuk membuat embedding text. Embedding text selanjutnya digunakan untuk menemukan embedding gambar serupa teratas yang kemudian ditampilkan di aplikasi. Jumlah hasil pencarian dan lebar gambar bisa dipilih menggunakan penggeser.

Demo

Demo aplikasi streamlit ditampilkan di bawah untuk kueri pencarian anjing pada kumpulan data gambar internet.

Disarankan untuk menjalankan streamlit dalam Wide Model yang bisa diakses dari menu pengaturan atas kanan.

Perbaikan

Aplikasi yang dikembangkan menggunakan CLIP terlatih untuk mencari gambar yang pas dengan kueri text yang dimasukkan. Tetapi, mungkin ada beberapa aplikasi khusus di mana CLIP terlatih tidak sesuai. Misalkan, menemukan merk mobil tertentu dari database banyak mobil merupakan tugas khusus. Untuk tugas semacam ini, kita perlu menyempurnakan CLIP pada kumpulan data khusus mobil berlabel. Posting blog ke-2 dalam seri ini akan menunjukkan cara menyempurnakan CLIP pada kumpulan data khusus domain.