<jupyter_start><jupyter_text>Image similarity search指定された画像と似ている画像を検索する。Step 1. Define AutoEncoder* 画像のencoder として AutoEncoder を定義する。Step 2. Lightning Module* encoder を pytorchlightning 用に定義する。Step 3. Training Encoder* encoder 訓練する。* 訓練済みモデルをファイルに保存する。Step 4. Embedding* 検索対象画像と検索画像を特徴ベクトルへ変換する。* 特徴ベクトルをファイルに保存する。Step 5. Faiss indexes* Step 4 で作成した検索対象画像の特徴ベクトルを検索するためのインデックスを生成する。* インデックスをファイルに保存する。Step 6. Image similarity search* Step 4 で作成した検索画像の特徴ベクトルを使用して、Step 5 のインデックスを検索する。----pytorch:* https://pytorch.org/pytorchligtning:* https://www.pytorchlightning.ai/faiss:* https://github.com/facebookresearch/faiss* https://faiss.ai/STL10:* https://pytorch.org/vision/stable/datasets.htmlAutoEndoder* https://qiita.com/MuAuan/items/a062d0c245c8f4836399  Setup for Notebook<jupyter_code># notebook runtime
import sys

runtime = 'local'
if 'google.colab' in sys.modules:
    runtime = 'colab'
elif _dh == ['/kaggle/working']:
    runtime = 'kaggle'
runtime
if runtime == 'colab':
    from google.colab import drive
    drive.mount('/content/drive')
if runtime == 'colab':
    home_path = '/content/drive/MyDrive/image_similarity_search'
else:
    home_path = '/home/jovyan/image_similarity_search'

nbs_path = f'{home_path}/nbs'
datasets_path = f'{home_path}/datasets'
models_path = f'{home_path}/models'
figs_path = f'{home_path}/figs'
logs_path = f'{home_path}/logs'
%cd {nbs_path}
!pip install -q pytorch_lightning
!pip install -q faiss-gpu<jupyter_output><empty_output><jupyter_text>Setup<jupyter_code>import os
import time
import numpy as np
from tqdm import tqdm_notebook as tqdm
import matplotlib.pyplot as plt
import pickle
import shutil
from PIL import Image
import glob
import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, random_split
from torchsummary import summary

import torchvision
from torchvision import transforms
import pytorch_lightning as pl
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import TensorBoardLogger
AVAIL_GPUS = min(1, torch.cuda.device_count())
BATCH_SIZE = 256 if AVAIL_GPUS else 64
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
pl.seed_everything(42)<jupyter_output><empty_output><jupyter_text>Dataset<jupyter_code>from torchvision.datasets import STL10
dataset_name = 'stl10'
def clear_datasets():
    path_cifar_10_batches_py = os.path.join(datasets_path, 'stl10_binary')
    if os.path.exists(path_cifar_10_batches_py):
        shutil.rmtree(path_cifar_10_batches_py)

    path_cifar_10_python = os.path.join(datasets_path, 'stl10_binary.tar.gz')
    if os.path.exists(path_cifar_10_python):
        os.remove(path_cifar_10_python)
#clear_datasets()
def train_transforms():
    return torchvision.transforms.Compose([
        torchvision.transforms.RandomCrop(32, padding=4),
        torchvision.transforms.RandomHorizontalFlip(),
        torchvision.transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
    ])
def test_transforms():
    return torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
        #cifar10_normalization(),
    ])
class DataModule(pl.LightningDataModule):

    def __init__(self, data_dir: str, batch_size: int=32):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.num_workers = int(os.cpu_count() / 2)
        self.train_transform = train_transforms()
        self.test_transform = test_transforms()
        self.classes = ['airplane', 'bird', 'car', 'cat', 'deer', 'dog', 'horse', 'monkey', 'ship', 'truck']
        self.num_classes = len(self.classes)
        self.dims = (3, 96, 96) # channels, width, height

    def prepare_data(self):
        # download
        STL10(self.data_dir, split='train', download=True)
        STL10(self.data_dir, split='test', download=True)

    def setup(self, stage=None): #train, val, testデータ分割
        # Assign train/val datasets for use in dataloaders
        datasets = STL10(self.data_dir, split='train', transform=self.train_transform)
        n_train = int(len(datasets) * 0.8)
        n_val = len(datasets) - n_train
        self.ds_train, self.ds_val = torch.utils.data.random_split(datasets, [n_train, n_val])
        self.ds_test = STL10(self.data_dir, split='test', transform=self.test_transform)

    def train_dataloader(self):
        return DataLoader(self.ds_train, shuffle=True, drop_last=True, batch_size=self.batch_size, num_workers=0)

    def val_dataloader(self):
        return DataLoader(self.ds_val, shuffle=False, batch_size=self.batch_size, num_workers=0)
 
    def test_dataloader(self):
        return DataLoader(self.ds_test, shuffle=False, batch_size=self.batch_size, num_workers=0)
datamodule = DataModule(datasets_path, BATCH_SIZE)
datamodule.prepare_data()
datamodule.setup()
# check data size
train_dataloader = iter(datamodule.train_dataloader())
images, labels = next(train_dataloader)
images.shape, labels.shape<jupyter_output><empty_output><jupyter_text>Step 1. Define AutoEncoder画像の特徴ベクトルを抽出するためAutoEncoder を定義する。<jupyter_code>class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.BatchNorm2d(64),

            nn.Conv2d(64, 256, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.BatchNorm2d(256),

            nn.Conv2d(256, 512, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.BatchNorm2d(512)
        )

    def forward(self, x):
        x = self.encoder(x)
        return x
encoder = Encoder()
encoder(torch.randn(32, 3, 96, 96)).shape
class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(in_channels = 512, out_channels = 256,
                                 kernel_size = 2, stride = 2, padding = 0),
            nn.ConvTranspose2d(in_channels = 256, out_channels = 16,
                                 kernel_size = 2, stride = 2, padding = 0),
            nn.ConvTranspose2d(in_channels = 16, out_channels = 3,
                                 kernel_size = 2, stride = 2)
        )

    def forward(self, x):
        x = self.decoder(x)
        return x
decoder = Decoder()
decoder(torch.randn(32, 512, 12, 12)).shape<jupyter_output><empty_output><jupyter_text>Step 2. Lightning Module<jupyter_code>encoder_name = 'autoencoder'
class LitAutoEncoder(pl.LightningModule):

    def __init__(self):
        super().__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()

    def forward(self, x):
        z = self.encoder(x)
        x_hat = self.decoder(z)
        return x_hat

    def training_step(self, batch, batch_idx):
        return self._common_step(batch, batch_idx, "train")

    def validation_step(self, batch, batch_idx):
        return self._common_step(batch, batch_idx, "val")

    def test_step(self, batch, batch_idx):
        return self._common_step(batch, batch_idx, "test")
    
    def _common_step(self, batch, batch_idx, stage: str):
        x, y = batch
        loss = F.mse_loss(x, self(x))
        self.log(f"{stage}_loss", loss, on_step=True)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer
# model
encoder = LitAutoEncoder()
encoder = encoder.to(device)  #for gpu
summary(encoder.encoder,(3, 96, 96))
summary(encoder.decoder,(512, 12, 12))
summary(encoder,(3, 96, 96))
print(encoder)<jupyter_output><empty_output><jupyter_text>Step 3. Training Encoder<jupyter_code>path_encoder = f'{models_path}/{dataset_name}_{encoder_name}.ckpt'
def clear_encoder():
    if os.path.exists(path_encoder):
        os.remove(path_encoder)
clear_encoder()
callbacks = [EarlyStopping(monitor="val_loss")]
trainer = pl.Trainer(gpus=1, callbacks=callbacks,
                logger=TensorBoardLogger(f'{logs_path}/lightning_logs/', name=f'{dataset_name}_{encoder_name}'))

if not os.path.exists(path_encoder):
    #trainer = pl.Trainer(max_epochs=10, gpus=1)
    trainer.fit(encoder, datamodule)
    trainer.save_checkpoint(path_encoder)
else:
    encoder = encoder.load_from_checkpoint(path_encoder)
encoder = encoder.to(device)  #for gpu
encoder.freeze()
encoder.eval()
# Test
# 0.013129102066159248
results = trainer.test(encoder, datamodule)
# functions to show an image
def imshow(img, file=None, title=None):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.detach().numpy()
    plt.figure(figsize=(20, 10))
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    if title:
        plt.title(title)
    if file:
        plt.savefig(file + '.png')
    plt.show()
# Original train images
dataiter = iter(datamodule.train_dataloader())
images, labels = dataiter.next()
imshow(torchvision.utils.make_grid(images, nrow=32), f'{figs_path}/{dataset_name}_{encoder_name}_train_original', title='train images')
print(' '.join('%5s' % datamodule.classes[labels[j]] for j in range(8)))
# Decoded train images
images_hat = encoder(images)
imshow(torchvision.utils.make_grid(images_hat, nrow=32), f'{figs_path}/{dataset_name}_{encoder_name}_train_decoded', title='train decoded images')
print(' '.join('%5s' % datamodule.classes[labels[j]] for j in range(8)))
# Original test images
dataiter = iter(datamodule.test_dataloader())
images, labels = dataiter.next()
imshow(torchvision.utils.make_grid(images, nrow=32), f'{figs_path}/{dataset_name}_{encoder_name}_test_original', title='test images')
print(' '.join('%5s' % datamodule.classes[labels[j]] for j in range(8)))
images_hat = encoder(images)
imshow(torchvision.utils.make_grid(images_hat, nrow=32), f'{figs_path}/{dataset_name}_{encoder_name}_test_decoded', title='test decoded images')
print(' '.join('%5s' % datamodule.classes[labels[j]] for j in range(8)))<jupyter_output><empty_output><jupyter_text>Step 4. Embedding検索対象画像と検索画像を特徴ベクトルへ変換する。<jupyter_code>path_embeded_train = f'{models_path}/{dataset_name}_{encoder_name}_embeded_train.pickle'
path_embeded_test = f'{models_path}/{dataset_name}_{encoder_name}_embeded_test.pickle'
def clear_embedding():
    if os.path.exists(path_embeded_train):
        os.remove(path_embeded_train)
    if os.path.exists(path_embeded_test):
        os.remove(path_embeded_test)
clear_embedding()
train_dataset = STL10(datasets_path, split='train', download=True)
test_dataset = STL10(datasets_path, split='test', download=True)
len(train_dataset), len(test_dataset)
# with preprocess
train_dataloader = DataLoader(STL10(datasets_path, split='train', download=True, transform=test_transforms()),
                               shuffle=False, batch_size=32, num_workers=0)
test_dataloader = DataLoader(STL10(datasets_path, split='test', download=True, transform=test_transforms()),
                               shuffle=False, batch_size=32, num_workers=0)
len(train_dataloader), len(test_dataloader)
def create_embedder_model(autoencoder):
    layers = list(autoencoder.encoder.children())
    #fv = nn.Sequential(nn.AdaptiveMaxPool2d(output_size=1))
    fv = nn.Sequential(nn.Flatten(), nn.Linear(in_features=512*12*12, out_features=512))
    model = nn.Sequential(*layers, *fv)
    return model
embedder = create_embedder_model(encoder)
embedder
img_random = torch.randn(32, 3, 96, 96)
img_emb = embedder(img_random)
img_emb.shape
def get_embeded_vector(embedder, dataloader):
    vector = []
    for i, (images, labels) in enumerate(tqdm(dataloader)):
        with torch.no_grad():
            v = embedder(images).squeeze().cpu()
        vector.extend(v.detach().numpy())

    return vector
if not os.path.exists(path_embeded_train):
    train_vectors = get_embeded_vector(embedder, train_dataloader)
    print(len(train_vectors), train_vectors[0].shape)
    with open(path_embeded_train, mode='wb') as f:
        pickle.dump(train_vectors, f)
if not os.path.exists(path_embeded_test):
    test_vectors = get_embeded_vector(embedder, test_dataloader)
    print(len(test_vectors))
    with open(path_embeded_test, mode='wb') as f:
        pickle.dump(test_vectors, f)<jupyter_output><empty_output><jupyter_text>Step 5. Faiss indexes検索用インデックスを生成する。<jupyter_code>import faiss
path_indexer = f'{models_path}/{dataset_name}_{encoder_name}_indexer.faiss'
def clear_indexer():
    if os.path.exists(path_indexer):
        os.remove(path_indexer)
clear_indexer()
with open(path_embeded_train, mode='rb') as f:
    train_vectors = np.array(pickle.load(f))

train_vectors.shape
class FlatIndexer(object):

    def __init__(self, vector_sz: int, nlist=10, path=None):
        if path and os.path.exists(path):
            index_cpu = faiss.read_index(path)
            self.indexer = faiss.index_cpu_to_all_gpus(index_cpu)
        else:
            #index_cpu = faiss.IndexFlatIP(vector_sz) # Not Work
            quantizer = faiss.IndexFlatL2(vector_sz)
            index_cpu = faiss.IndexIVFFlat(quantizer, vector_sz, nlist, faiss.METRIC_L2)
            res = faiss.StandardGpuResources()
            self.indexer = faiss.index_cpu_to_gpu(res, 0, index_cpu)

    def index_data(self, vectors):
        self.indexer.train(vectors)
        self.indexer.add(vectors)

    def search_knn(self, query_vectors: np.array, top_docs: int):
        scores, indexes = self.indexer.search(query_vectors, top_docs)
        return scores, indexes

    def save_index(self, path):
        index_cpu = faiss.index_gpu_to_cpu(self.indexer)
        faiss.write_index(index_cpu, path)
if not os.path.exists(path_indexer):
    indexer = FlatIndexer(512, nlist=20)
    indexer.index_data(train_vectors)
    indexer.save_index(path_indexer)
else:
    indexer = FlatIndexer(512, path_indexer)<jupyter_output><empty_output><jupyter_text>Step 6. Image similarity search<jupyter_code>with open(path_embeded_test, mode='rb') as f:
    test_vectors = np.array(pickle.load(f))

test_vectors.shape
scores, indexes = indexer.search_knn(test_vectors, 10)
scores.shape, indexes.shape
for i in range(0, 20):
    print(i, indexes[i])
scores[0]
def show_search_images(train_dataset, test_dataset, indexes, i, dataset_name, encoder_name):
    " train_dataset の中から test_dataset で指定された画像に似ている画像を検索する "
    fig, axes = plt.subplots(1, 11, figsize=(15,5))
    test_image, test_label = test_dataset[i]
    axes[0].imshow(test_image)
    axes[0].set_xticks([])
    axes[0].set_yticks([])
    axes[0].set_title(f'Q[{i}]')

    for col, idx in enumerate(indexes[i]):
        img, label = train_dataset[idx]
        axes[col+1].set_title(f'A[{idx}]')
        axes[col+1].imshow(img)
        axes[col+1].set_xticks([])
        axes[col+1].set_yticks([])
    plt.show()
    fig.savefig(f'{figs_path}/{dataset_name}_{encoder_name}_search_images_{i}')
for i in range(10):
    show_search_images(train_dataset, test_dataset, indexes, i, dataset_name, encoder_name)
def merge_search_images(figs_path, dataset_name, encoder_name):
    files = glob.glob(f"{figs_path}/{dataset_name}_{encoder_name}_search_images_*.png")
    images = None
    for file in sorted(files):
        im = np.array(Image.open(file).convert('RGB'))
        h, w, c = im.shape
        im = im[120:h-130, 110:w-80, :] # trim
        if images is None:
            images = im
        else:
            images = np.append(images, im, axis=0)
    img = Image.fromarray(images)
    img.save(f"{figs_path}/{dataset_name}_{encoder_name}_search_images.png")
    return img
img = merge_search_images(figs_path, dataset_name, encoder_name)
img
# !nb-clean clean stl10_autoencoder.ipynb<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12464607059711823, 'fraction_numerical': 0.021896432391618953, 'mean_word_length': 4.182262797522008, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 37, 'https://': 6, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '338580', 'license': 'MIT', 'n_tokens_mistral': 6033, 'n_tokens_neox': 5728, 'n_words': 1009, 'path': 'nbs/stl10_autoencoder.ipynb', 'repo_name': 'benymd/image_search', 'size': '15893'}