Upload 21 files

README.md

@@ -1,11 +1,129 @@
 ---
-title: Img2art Search
-emoji: 🔥
-colorFrom: yellow
-colorTo: blue
-sdk: docker
-pinned: false
-license: mit
+title: img2art-search
+app_file: app.py
+sdk: gradio
+sdk_version: 4.37.2
 ---
+# Image-to-Art Search 🔍
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+"<b>Find real artwork that looks like your images</b>"
+
+This project fine-tunes a Vision Transformer (ViT) model, pre-trained with "google/vit-base-patch32-224-in21k" weights and fine tuned with the style of [ArtButMakeItSports](https://www.instagram.com/artbutmakeitsports/), to perform image-to-art search across 81k artworks made available by [WikiArt](https://wikiart.org/).
+
+![horse](examples/horse.jpg)
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Installation](#installation)
+- [How it works](#how-it-works)
+- [Dataset](#dataset)
+- [Training](#training)
+- [Inference](#inference)
+- [Contributing](#contributing)
+- [License](#license)
+
+## Overview
+
+This project leverages the Vision Transformer (ViT) model architecture for the task of image-to-art search. By fine-tuning the pre-trained ViT model on a custom dataset derived from the Instagram account [ArtButMakeItSports](https://www.instagram.com/artbutmakeitsports/), we aim to create a model capable of matching images (but not only) to corresponding artworks, being able to search for any of the images on [WikiArt](https://wikiart.org/).
+
+## Installation
+
+1. Clone the repository:
+```sh
+git clone https://github.com/brunorosilva/img2art-search.git
+cd img2art-search
+```
+
+2. Install poetry:
+```sh
+pip install poetry
+```
+
+3. Install using poetry:
+```sh
+poetry install
+```
+
+## How it works
+
+### Dataset Preparation
+
+1. Download images from the [ArtButMakeItSports](https://www.instagram.com/artbutmakeitsports/) Instagram account.
+2. Organize the images into appropriate directories for training and validation.
+3. Get a fine tuned model
+4. Create the gallery using WikiArt
+
+### Training
+
+Fine-tune the ViT model:
+```sh
+make train
+```
+
+### Inference via Gradio
+
+Perform image-to-art search using the fine-tuned model:
+```sh
+make viz
+```
+
+### Recreate the wikiart gallery
+```sh
+make wikiart
+```
+
+### Create new gallery
+
+If you want to index new images to search, use:
+```sh
+poetry run python main.py gallery --gallery_path <your_path>
+```
+
+## Dataset
+
+The dataset derives from 1k images from the Instagram account [ArtButMakeItSports](https://www.instagram.com/artbutmakeitsports/). Images are downloaded and split into training, validation and test sets. Each image is paired with its corresponding artwork for training purposes, if you want this dataset just ask me stating your usage.
+
+WikiArt is indexed using the same process, except that there's no expected result. So each artwork is mapped to itself and the model is used as a feature extractor and the gallery embeddings are saved as a numpy file (will be changed to chromadb in the future).
+
+## Training
+
+The training script fine-tunes the ViT model on the prepared dataset. Key steps include:
+
+1. Loading the pre-trained "google/vit-base-patch32-224-in21k" weights.
+2. Preparing the dataset and data loaders.
+3. Fine-tuning the model using a custom training loop.
+4. Saving the model to the models folder
+
+## Interface
+
+The recommended method to get models is to use [gradio](https://www.gradio.app/) as an interface by running `make viz`. This will open a server and you can use some image you want to search or even use your webcam to get top 4 search models.
+
+### Examples
+Search for contextual similarity
+![field](examples/field.jpg)
+
+Search for shapes similarity
+![basket](examples/basketball.jpg)
+
+Search for expression similarity (yep, that's me)
+![serious_face](examples/serious_face.jpg)
+
+Search for pose similarity
+![lawyer](examples/lawyer.jpg)
+
+Search for an object
+![horse](examples/horse.jpg)
+
+
+
+## Contributing
+There are three topics I'd appreciate help with:
+1. Increasing the gallery by embedding new painting datasets, the current one has 81k artworks because I just got a ready to go dataset, but the complete WikiArt catalog alone has 250k+ artworks, so I really want to up this number to a least 300k in the near future;
+2. Porting the encoding and search to a vector db, preferably chromadb;
+3. Open issues with how this could be improved, new ideas will be considered.
+
+## License
+The source code for the site is licensed under the MIT license, which you can find in the MIT-LICENSE.txt file.
+
+All graphical assets are licensed under the Creative Commons Attribution 3.0 Unported License.

img2art_search/constants.py

@@ -0,0 +1 @@
+BASE_PATH = "data/artmakeitsports"

img2art_search/data/data.py

@@ -0,0 +1,28 @@
+import os
+
+import numpy as np
+
+from img2art_search.constants import BASE_PATH
+
+
+def get_data_from_local() -> np.ndarray:
+    left_or_top_data = [
+        f"{BASE_PATH}/splits/left_or_top/{fn}"
+        for fn in os.listdir(f"{BASE_PATH}/splits/left_or_top")
+    ]
+
+    x = np.array(left_or_top_data)
+    y = np.array([ex.replace("left_or_top", "right_or_bottom") for ex in x])
+
+    data = np.array([x, y])
+    return data
+
+
+def split_train_val_test(data: np.ndarray, test_size: float, val_size: float) -> tuple:
+    train_size = 1 - test_size - val_size
+    SPLIT = int(data.shape[1] * train_size)
+    TEST_SPLIT = SPLIT + int(data.shape[1] * test_size)
+    train = data[:, :SPLIT]
+    validation = data[:, SPLIT:TEST_SPLIT]
+    test = data[:, TEST_SPLIT:]
+    return train, validation, test

img2art_search/data/dataset.py

@@ -0,0 +1,24 @@
+import numpy as np
+from PIL import Image
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+
+class ImageRetrievalDataset(Dataset):
+    def __init__(self, data: np.ndarray, transform: transforms.Compose) -> None:
+        self.data = data
+        self.transform = transform
+
+    def __len__(self) -> int:
+        return len(self.data[0])
+
+    def __getitem__(self, idx: int) -> tuple:
+        input_path, label_path = self.data.T[idx]
+        input_image = Image.open(input_path).convert("RGB")
+        label_image = Image.open(label_path).convert("RGB")
+
+        # if self.transform:
+        input_image = self.transform(input_image)
+        label_image = self.transform(label_image)
+
+        return input_image, label_image, input_path, label_path

img2art_search/data/transforms.py

@@ -0,0 +1,19 @@
+from torchvision import transforms
+
+transform = transforms.Compose(
+    [
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ]
+)
+
+
+inversetransform = transforms.Compose(
+    [
+        transforms.Normalize(
+            mean=[0.0, 0.0, 0.0], std=[1 / 0.229, 1 / 0.224, 1 / 0.225]
+        ),
+        transforms.Normalize(mean=[-0.485, -0.456, -0.406], std=[1.0, 1.0, 1.0]),
+    ]
+)

img2art_search/losses/contrastiveloss.py

@@ -0,0 +1,12 @@
+import torch
+
+
+class ContrastiveLoss(torch.nn.Module):
+    def __init__(self, margin=1.0):
+        super(ContrastiveLoss, self).__init__()
+        self.margin = margin
+
+    def forward(self, output1, output2):  # noqa
+        euclidean_distance = torch.nn.functional.pairwise_distance(output1, output2)
+        loss = torch.mean(torch.pow(euclidean_distance, 2))
+        return loss

img2art_search/models/compute_embeddings.py

@@ -0,0 +1,120 @@
+import base64
+import os
+from io import BytesIO
+
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from img2art_search.data.dataset import ImageRetrievalDataset
+from img2art_search.data.transforms import transform
+from img2art_search.models.model import ViTImageSearchModel
+from img2art_search.utils import (
+    get_or_create_pinecone_index,
+    get_pinecone_client,
+    inverse_transform_img,
+)
+
+
+def extract_embedding(image_data_batch, fine_tuned_model):
+    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    image_data_batch = image_data_batch.to(DEVICE)
+    with torch.no_grad():
+        embeddings = fine_tuned_model(image_data_batch).cpu().numpy()
+    return embeddings
+
+
+def load_fine_tuned_model():
+    fine_tuned_model = ViTImageSearchModel()
+    fine_tuned_model.load_state_dict(torch.load("models/model.pth"))
+    fine_tuned_model.eval()
+    return fine_tuned_model
+
+
+def create_gallery(
+    img_dataset: ImageRetrievalDataset,
+    fine_tuned_model: ViTImageSearchModel,
+    save: bool = True,
+) -> list:
+    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    batch_size = 4
+    fine_tuned_model.to(DEVICE)
+    gallery_embeddings = []
+    gallery_dataloader = DataLoader(
+        img_dataset, batch_size=batch_size, num_workers=1, shuffle=False
+    )
+
+    pc = get_pinecone_client()
+    gallery_index = get_or_create_pinecone_index(pc)
+    try:
+        count = 0
+        for img_data, _, img_name, _ in tqdm(gallery_dataloader):
+            data_objects = []
+            batch_embedding = extract_embedding(img_data, fine_tuned_model)
+            gallery_embeddings.append(batch_embedding)
+            for idx, embedding in enumerate(batch_embedding):
+                image = Image.fromarray(
+                    inverse_transform_img(img_data[idx]).numpy().astype("uint8"), "RGB"
+                )
+                buffered = BytesIO()
+                image.save(buffered, format="JPEG")
+                img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+                data_objects.append(
+                    {
+                        "id": str(count),
+                        "values": embedding.tolist(),
+                        "metadata": {
+                            "image": img_base64,
+                            "name": img_name[idx]
+                            .split("/")[-1]
+                            .replace(".jpg", "")
+                            .replace(".jpeg", "")
+                            .replace(".png", "")
+                            .replace(".JPG", "")
+                            .replace(".JPEG", "")
+                            .replace("-", " ")
+                            .replace("_", " - ")
+                            .title(),
+                        },
+                    }
+                )
+                count += 1
+            gallery_index.upsert(vectors=data_objects)
+    except Exception as e:
+        print(f"Error creating gallery: {e}")
+
+    if save:
+        np.save("models/embeddings", gallery_embeddings)
+    return gallery_embeddings
+
+
+def search_image(query_image_path: str, k: int = 4) -> tuple:
+    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    fine_tuned_model = load_fine_tuned_model()
+    fine_tuned_model.to(DEVICE)
+    query_embedding = extract_embedding(query_image_path, fine_tuned_model)
+    pc = get_pinecone_client()
+    index = get_or_create_pinecone_index(pc)
+    response = index.query(
+        vector=[query_embedding.tolist()[0]], top_k=k, include_metadata=True
+    )
+    distances = []
+    results = []
+    for obj in response["matches"]:
+        result = base64.b64decode(obj.metadata["image"])
+        results.append(result)
+        distances.append(
+            str(round(obj["score"], 2) * 100) + " " + str(obj.metadata["name"])
+        )
+
+    return results, distances
+
+
+def create_gallery_embeddings(folder: str) -> None:
+    x = np.array([f"{folder}/{file}" for file in os.listdir(folder)])
+    gallery_data = np.array([x, x])
+    gallery_dataset = ImageRetrievalDataset(gallery_data, transform=transform)
+    fine_tuned_model = load_fine_tuned_model()
+    create_gallery(gallery_dataset, fine_tuned_model)

img2art_search/models/model.py

@@ -0,0 +1,13 @@
+from torch import nn
+from transformers import ViTModel
+
+
+class ViTImageSearchModel(nn.Module):
+    def __init__(self, pretrained_model_name="google/vit-base-patch32-224-in21k"):
+        super(ViTImageSearchModel, self).__init__()
+        self.vit = ViTModel.from_pretrained(pretrained_model_name)
+
+    def forward(self, x):  # noqa
+        outputs = self.vit(pixel_values=x)
+        cls_hidden_state = outputs.last_hidden_state[:, 0, :]
+        return cls_hidden_state

img2art_search/models/predict.py

@@ -0,0 +1,37 @@
+import os
+from io import BytesIO
+
+import numpy as np
+import torch
+from PIL import Image
+
+from img2art_search.data.dataset import ImageRetrievalDataset
+from img2art_search.data.transforms import transform
+from img2art_search.models.compute_embeddings import search_image
+
+
+def predict(img: Image.Image) -> list:
+    tmp_img_path = "tmp_img.png"
+    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    if img:
+        img.save(tmp_img_path)
+        pred_img = np.array([[tmp_img_path], [tmp_img_path]])
+        pred_dataset = ImageRetrievalDataset(pred_img, transform=transform)
+        pred_image_data = pred_dataset[0][0].unsqueeze(0).to(DEVICE)
+        indices, distances = search_image(pred_image_data)
+        results = []
+        for index, distance in zip(indices, distances):
+            buffered = BytesIO(index)
+            image = Image.open(buffered)
+            decoded_image_array = np.array(image)
+
+            results.append(
+                (
+                    Image.fromarray(decoded_image_array),
+                    str(distance),
+                )
+            )
+        os.remove(tmp_img_path)
+        return results
+    else:
+        return []

img2art_search/models/train.py

@@ -0,0 +1,75 @@
+import numpy as np
+import torch
+from torch.optim import Adam
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+
+from img2art_search.data.data import get_data_from_local, split_train_val_test
+from img2art_search.data.dataset import ImageRetrievalDataset
+from img2art_search.data.transforms import transform
+from img2art_search.losses.contrastiveloss import ContrastiveLoss
+from img2art_search.models.model import ViTImageSearchModel
+
+
+def fine_tune_vit(epochs: int, batch_size: int) -> None:
+    data = get_data_from_local()
+    train_data, val_data, test_data = split_train_val_test(data, 0.2, 0.1)
+    np.save("models/test_data", test_data)
+    train_dataset = ImageRetrievalDataset(train_data, transform=transform)
+    val_dataset = ImageRetrievalDataset(val_data, transform=transform)
+
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
+
+    model = ViTImageSearchModel()
+
+    # logs
+    log_dir = "./logs/"
+    writer = SummaryWriter(log_dir=log_dir)
+
+    # params
+    criterion = ContrastiveLoss()
+    optimizer = Adam(model.parameters(), lr=1e-4)
+    epochs = epochs
+
+    for epoch in range(epochs):
+        model.train()
+        total_loss = 0
+
+        for batch_idx, batch in enumerate(train_loader):
+            inputs, labels = batch
+            optimizer.zero_grad()
+
+            input_embeddings = model(inputs)
+            label_embeddings = model(labels)
+
+            loss = criterion(input_embeddings, label_embeddings)
+
+            loss.backward()
+            optimizer.step()
+
+            total_loss += loss.item()
+            writer.add_scalar(
+                "Train Loss", loss.item(), epoch * len(train_loader) + batch_idx
+            )
+
+        avg_train_loss = total_loss / len(train_loader)
+        writer.add_scalar("Average Train Loss", avg_train_loss, epoch)
+
+        print(f"Epoch [{epoch + 1}/{epochs}], Loss: {total_loss / len(train_loader)}")
+
+        model.eval()
+        with torch.no_grad():
+            val_loss = 0
+            for batch_idx, batch in enumerate(val_loader):
+                inputs, labels = batch
+                input_embeddings = model(inputs)
+                label_embeddings = model(labels)
+
+                loss = criterion(input_embeddings, label_embeddings)
+                val_loss += loss.item()
+            avg_val_loss = val_loss / len(val_loader)
+            writer.add_scalar("Validation Loss", avg_val_loss, epoch)
+            print(f"Validation Loss: {val_loss / len(val_loader)}")
+
+    torch.save(model.state_dict(), "models/model.pth")

img2art_search/utils.py

@@ -0,0 +1,36 @@
+import os
+from typing import Any
+
+import torch
+from pinecone import Pinecone, ServerlessSpec
+
+from img2art_search.data.transforms import inversetransform
+
+pinecone_api_key = os.environ["PINECONE_API_KEY"]
+
+def inverse_transform_img(img: torch.Tensor) -> torch.Tensor:
+    inv_tensor = inversetransform(img)
+    tensor_image = (inv_tensor * 255).byte()
+    return tensor_image.permute(1, 2, 0)
+
+
+def get_pinecone_client() -> Pinecone:
+    pc = Pinecone(api_key=pinecone_api_key)
+    return pc
+
+
+def get_or_create_pinecone_index(
+    pc: Pinecone, index_name: str = "img2art-search", embeddings_dim: int = 768
+) -> Any:
+    indexes_names = [index.name for index in pc.list_indexes()]
+    if index_name not in indexes_names:
+        pc.create_index(
+            name=index_name,
+            dimension=embeddings_dim,
+            metric="cosine",
+            spec=ServerlessSpec(cloud="aws", region="us-east-1"),
+        )
+
+    index = pc.Index(index_name)
+
+    return index

main.py

@@ -0,0 +1,48 @@
+from img2art_search.models.predict import predict
+from img2art_search.models.train import fine_tune_vit
+from img2art_search.models.compute_embeddings import create_gallery_embeddings
+import gradio as gr
+import argparse
+
+def make_interface():
+    interface = gr.Interface(
+        fn=predict,
+        inputs=gr.Image(type="pil"),
+        outputs=gr.Gallery(label="Most similar images", height=256 * 3),
+        # live=True,
+    )
+    interface.launch(share=False)
+
+def train(epochs, batch_size):
+    fine_tune_vit(epochs, batch_size)
+
+def create_gallery(gallery_path):
+    create_gallery_embeddings(gallery_path)
+
+def main():
+    parser = argparse.ArgumentParser(description="Train or infer the ViT model for image-to-art search.")
+    subparsers = parser.add_subparsers(dest="command")
+
+    # Subparser for training
+    train_parser = subparsers.add_parser("train", help="Fine-tune the ViT model")
+    train_parser.add_argument("--epochs", type=int, default=50, help="Number of training epochs")
+    train_parser.add_argument("--batch_size", type=int, default=32, help="Batch size for training")
+
+    # Subparser for inference
+    _ = subparsers.add_parser("interface", help="Perform image-to-art search using the fine-tuned model")
+
+    create_gallery_parser = subparsers.add_parser("gallery", help="Create new gallery from a path")
+    create_gallery_parser.add_argument("--gallery_path", type=str, default="data/wikiart")
+    args = parser.parse_args()
+
+    if args.command == "train":
+        train(args.epochs, args.batch_size)
+    elif args.command == "interface":
+        make_interface()
+    elif args.command == "gallery":
+        create_gallery(args.gallery_path)
+    else:
+        parser.print_help()
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -0,0 +1,38 @@
+[tool.poetry]
+name = "img2art-search"
+version = "0.1.0"
+description = ""
+authors = ["brunorosilva <[email protected]>"]
+readme = "README.md"
+packages = [{include = "img2art_search"}]
+package-mode = false
+
+[tool.poetry.dependencies]
+python = "^3.10"
+torch = "^2.3.1"
+torchvision = "^0.18.1"
+tqdm = "^4.66.4"
+pandas = "^2.2.2"
+scipy = "^1.14.0"
+numpy = "^2.0.0"
+transformers = "^4.41.2"
+tensorboard = "^2.17.0"
+scikit-learn = ">=0.0.0"
+matplotlib = "^3.9.0"
+gradio = ">=0.0.0"
+opencv-python = "^4.10.0.84"
+pinecone = "^6.0.1"
+
+
+[tool.poetry.dev-dependencies]
+black = "^24.4.2"
+vulture = "^2.11"
+mypy = "^1.10.1"
+flake8 = "^7.1.0"
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[[tool.mypy.overrides]]
+module = ["torchvision.*", "transformers.*", "pinecone.*", "tqdm.*"]
+follow_untyped_imports = true

requirements.txt

@@ -0,0 +1,11 @@
+torch
+torchvision
+tqdm
+pandas
+scipy
+numpy
+transformers
+tensorboard
+scikit-learn
+matplotlib
+gradio

setup.cfg

@@ -0,0 +1,10 @@
+[flake8]
+max-line-length = 88
+ignore = E122,E123,E126,E127,E128,E203,E221,E241,E731,E722,W503
+exclude = tests,.git,__init__.py
+
+# [mypy]
+# ignore_missing_imports = True
+
+[bdist_wheel]
+universal=1

setup.py

@@ -0,0 +1,36 @@
+import io
+import os
+import re
+
+from setuptools import find_packages
+from setuptools import setup
+
+
+def read(filename):
+    filename = os.path.join(os.path.dirname(__file__), filename)
+    text_type = type("")
+    with io.open(filename, mode="r", encoding="utf-8") as fd:
+        return re.sub(text_type(r":[a-z]+:`~?(.*?)`"), text_type(r"``\1``"), fd.read())
+
+
+requirements = [
+    # use environment.yml
+]
+
+
+setup(
+    name="img2art_search",
+    version="0.0.1",
+    url="https://github.com/brunorosilva/img2art-search",
+    author="Bruno Chicelli",
+    author_email="[email protected]",
+    description="Short description",
+    long_description=read("README.rst"),
+    packages=find_packages(exclude=("tests",)),
+    entry_points={"console_scripts": ["img2art_search=img2art_search.cli:cli"]},
+    install_requires=requirements,
+    classifiers=[
+        "Programming Language :: Python",
+        "Programming Language :: Python :: 3.6",
+    ],
+)