Add application file

app.py

@@ -0,0 +1 @@
+from fis.app import app  # noqa: F401

fis/app/app.py

@@ -0,0 +1,39 @@
+import os
+from typing import List
+
+import gradio as gr
+import numpy as np
+from datasets import load_dataset
+from PIL.Image import Image as Img
+
+from fis.feature_extraction.pipeline.pipeline import factory
+from fis.utils.constants import ORGANISATION
+from fis.utils.s3 import read_image_from_s3
+
+# Ugly fix of "OMP: Error #15: Initializing libomp.a, but found libiomp5.dylib already initialized."
+os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
+
+
+PIPELINE_NAME = "dummy_swin_pipe"
+
+pipeline = factory.get(PIPELINE_NAME)
+
+DATASET_PATH = os.path.join(ORGANISATION, "dummy_swin_pipe_debug")
+dataset = load_dataset(path=DATASET_PATH, split="train")
+dataset.add_faiss_index(column="embedding")
+
+
+def find_most_similar(image: np.ndarray) -> List[Img]:
+    image_embeddings = pipeline.encode(image)[0]
+
+    scores, samples = dataset.get_nearest_examples("embedding", image_embeddings, k=5)
+
+    images = []
+    for image_path in samples["path"]:
+        image = read_image_from_s3(image_path)
+        images.append(image)
+
+    return images
+
+
+gr.Interface(fn=find_most_similar, inputs="image", outputs=["image" for i in range(5)]).launch()

fis/feature_extraction/detection/base.py

@@ -0,0 +1,18 @@
+from abc import ABC, abstractmethod
+
+from PIL import Image
+
+
+class BaseDetector(ABC):
+    """Base class for detection models."""
+
+    @abstractmethod
+    def __call__(self, image: Image) -> None:
+        """Get embeddings from an image.
+
+        Args:
+            image: Image to encode
+
+        Returns:
+            Embedding
+        """

fis/feature_extraction/detection/dummy.py

@@ -0,0 +1,23 @@
+from typing import List, Tuple
+
+from PIL import Image
+
+from fis.feature_extraction.detection.base import BaseDetector
+
+
+class DummyDetector(BaseDetector):
+    """Dummy detection model."""
+
+    def __call__(self, image: Image) -> List[Tuple[int]]:
+        """Return a bounding box with the same size as the image.
+
+        Args:
+            image: Image
+
+        Returns:
+            Dummy bounding box the same size as the image
+        """
+        x_min, y_min = 0, 0
+        x_max, y_max = image.size
+
+        return [(x_min, y_min, x_max, y_max)]

fis/feature_extraction/embedding/base.py

@@ -0,0 +1,18 @@
+from abc import ABC, abstractmethod
+
+from PIL import Image
+
+
+class BaseEncoder(ABC):
+    """Base class for encoders."""
+
+    @abstractmethod
+    def __call__(self, image: Image) -> None:
+        """Get embeddings from an image.
+
+        Args:
+            image: Image to encode
+
+        Returns:
+            Embedding
+        """

fis/feature_extraction/embedding/timm.py

@@ -0,0 +1,63 @@
+from typing import Callable, Tuple
+
+import timm
+import torch
+from PIL import Image
+from timm.data import resolve_data_config
+from timm.data.transforms_factory import create_transform
+
+from fis.feature_extraction.embedding.base import BaseEncoder
+
+
+class TimmModel(BaseEncoder):
+    """Base class for timm models."""
+
+    def __init__(self, model_name: str) -> None:
+        """Instantiate the model class.
+
+        Args:
+            model_name: Name of the model in the timm library.
+        """
+        model, transform = self._creat_timm_model(model_name=model_name)
+
+        self._model_name = model_name
+        self._model = model
+        self._transform = transform
+
+    @property
+    def model_name(self) -> str:
+        return self._model_name
+
+    @staticmethod
+    def _creat_timm_model(model_name: str) -> Tuple[torch.nn.Module, Callable]:
+        """Create a model and its assitiated configuration.
+
+        Args:
+            model_name: Name of the model in the timm library.
+
+        Returns:
+            model and transformation function for input images.
+        """
+        model = timm.create_model(model_name=model_name, pretrained=True, num_classes=0)
+        model.eval()
+
+        config = resolve_data_config({}, model=model)
+        transform = create_transform(**config)
+
+        return model, transform
+
+    def __call__(self, image: Image) -> torch.Tensor:
+        """Get embeddings from an image.
+
+        Args:
+            image: Image to encode
+
+        Returns:
+            Embedding
+        """
+        tensor = self._transform(image).unsqueeze(0)  # transform and add batch dimension
+
+        with torch.no_grad():
+            embedding = self._model(tensor)
+
+        return embedding.numpy()[0]

fis/feature_extraction/pipeline/base.py

@@ -0,0 +1,85 @@
+from typing import Any, List
+
+import numpy as np
+import torch
+from PIL import Image
+from PIL.Image import Image as Img
+
+from fis.feature_extraction.detection.base import BaseDetector
+from fis.feature_extraction.embedding.base import BaseEncoder
+
+
+class EncodingPipeline:
+    """Apply the detection and embedding models to an image."""
+
+    def __init__(self, name: str, detection_model: BaseDetector, embedding_model: BaseEncoder) -> None:
+        """Initialize the encoding pipeline.
+
+        Args:
+            name: Name of the pipeline.
+            detection_model: Model used to detect the fashion items in the images.
+            embedding_model: Model used to generate embeddings for each detected item.
+        """
+        self._name = name
+        self._detection_model = detection_model
+        self._embedding_model = embedding_model
+
+    def encode(self, image: str) -> List[torch.Tensor]:
+        """Encode each item from an image into a embedding.
+
+        Args:
+            image: path to the image.
+
+        Returns:
+            Embeddings for each detected item in the image.
+        """
+        image = self._load_images(image)
+        bboxes = self._detection_model(image)
+        items = self._crop_images(image, bboxes)
+
+        embeddings = []
+        for item in items:
+            embedding = self._embedding_model(item)
+            embeddings.append(embedding)
+
+        return embeddings
+
+    def _load_images(self, image: Any) -> Img:
+        """Read an image from disk.
+
+        Args:
+            image: Path to the image on disk.
+
+        Raises:
+            TypeError: if the type of image is incorrect.
+
+        Returns:
+            PIL Image.
+        """
+        if isinstance(image, Img):
+            pass
+        elif isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        elif isinstance(image, str):
+            image = Image.open(image)
+        else:
+            raise TypeError(f"Unknown type for image: {type(image)}")
+
+        return image
+
+    def _crop_images(self, image, bboxes) -> List[Img]:
+        """Crop an image based on bounding boxes.
+
+        Args:
+            image: Image to crop items from.
+            bboxes: Bounding box containing an item.
+
+        Returns:
+            List of cropped images.
+        """
+        items = []
+        for bbox in bboxes:
+            cropped_image = image.crop(bbox)
+            items.append(cropped_image)
+
+        return items

fis/feature_extraction/pipeline/factory.py

@@ -0,0 +1,51 @@
+from fis.feature_extraction.detection.base import BaseDetector
+from fis.feature_extraction.embedding.base import BaseEncoder
+from fis.feature_extraction.pipeline.base import EncodingPipeline
+
+
+class PipelineFactory:
+    """Factory method for encoding pipelines.
+
+    Example use:
+        >>> from fis.feature_extraction.pipeline.factory import PipelineFactory
+        >>> factory = PipelineFactory()
+        >>> factory.register_pipeline(
+        ...     name="example_pipeline",
+        ...     detection_model=BaseDetector(),
+        ...     embedding_model=BaseEncoder()
+        ... )
+        >>> pipeline = factory.get('example_pipeline')
+    """
+
+    def __init__(self):
+        """Instantiate factory object."""
+        self._pipelines = {}
+
+    def register_pipeline(self, name: str, detection_model: BaseDetector, embedding_model: BaseEncoder) -> None:
+        """Register a new pipeline to the factory.
+
+        Args:
+            name: Name of the pipeline to create.
+            detection_model: Instance of a BaseDetector object.
+            embedding_model: Instance of a BaseEncoder object.
+        """
+        pipeline = EncodingPipeline(name=name, detection_model=detection_model, embedding_model=embedding_model)
+        self._pipelines[name] = pipeline
+
+    def get(self, name: str) -> EncodingPipeline:
+        """Get a pipeline from its name.
+
+        Args:
+            name: Name of the pipeline to get.
+
+        Raises:
+            ValueError: If no pipeline has been registered with the given name.
+
+        Returns:
+            Encoding pipeline.
+        """
+        pipeline = self._pipelines.get(name)
+        if not pipeline:
+            raise ValueError(name)
+
+        return pipeline

fis/feature_extraction/pipeline/pipeline.py

@@ -0,0 +1,11 @@
+from fis.feature_extraction.detection.dummy import DummyDetector
+from fis.feature_extraction.embedding.timm import TimmModel
+from fis.feature_extraction.pipeline.factory import PipelineFactory
+
+factory = PipelineFactory()
+
+factory.register_pipeline(
+    name="dummy_swin_pipe",
+    detection_model=DummyDetector(),
+    embedding_model=TimmModel(model_name="swinv2_base_window8_256"),
+)

fis/feature_extraction/run.py

@@ -0,0 +1,51 @@
+import pandas as pd
+from datasets import Dataset
+from tqdm import tqdm
+
+from fis.feature_extraction.pipeline.pipeline import factory
+from fis.utils.constants import ORGANISATION
+from fis.utils.s3 import list_images_from_bucket, read_image_from_s3
+
+
+def make_dataset(pipeline_name: str) -> Dataset:
+    print("Listing images from S3...")
+    images = list_images_from_bucket()
+    images = images[:100000]
+    print(f"{len(images)} images to process.")
+
+    pipeline = factory.get(pipeline_name)
+    data = []
+
+    print("Encoding images...")
+    for image_name in tqdm(images):
+        image = read_image_from_s3(image_name)
+        embeddings = pipeline.encode(image)
+
+        for embedding in embeddings:
+            image_data = {
+                "path": image_name,
+                "embedding": embedding,
+            }
+
+            data.append(image_data)
+
+    df = pd.DataFrame(data)
+    dataset = Dataset.from_pandas(df)
+
+    return dataset
+
+
+def upload_dataset(dataset: Dataset, pipeline_name: str) -> None:
+    print("Uploading dataset...")
+    repo_id = "{}/{}".format(ORGANISATION, pipeline_name)
+    dataset.push_to_hub(repo_id=repo_id)
+
+
+def main():
+    pipeline_name = "dummy_swin_pipe"
+    dataset = make_dataset(pipeline_name=pipeline_name)
+    upload_dataset(dataset=dataset, pipeline_name=pipeline_name)
+
+
+if __name__ == "__main__":
+    main()

fis/similarity_search/milvus/collection.py

@@ -0,0 +1,58 @@
+from pymilvus import (
+    Collection,
+    CollectionSchema,
+    DataType,
+    FieldSchema,
+    connections,
+    utility,
+)
+
+connections.connect(host="127.0.0.1", port="19530")
+
+
+def create_milvus_collection(collection_name: str, dim: int) -> Collection:
+    """Create a Milvus collection.
+
+    Inspired by https://github.com/milvus-io/bootcamp/blob/master/solutions/reverse_image_search/1_build_image_search_engine.ipynb
+
+    Args:
+        collection_name: name of the Milvus collection
+        dim: number of dimentions
+
+    Returns:
+        Milvus collection
+    """
+    if utility.has_collection(collection_name):
+        utility.drop_collection(collection_name)
+
+    fields = [
+        FieldSchema(
+            name="id",
+            dtype=DataType.INT64,
+            descrition="ids",
+            is_primary=True,
+            auto_id=False,
+        ),
+        FieldSchema(
+            name="path",
+            dtype=DataType.VARCHAR,
+            descrition="path to image",
+            max_length=500,
+            # is_primary=True,
+            # auto_id=False,
+        ),
+        FieldSchema(
+            name="embedding",
+            dtype=DataType.FLOAT_VECTOR,
+            descrition="image embedding vectors",
+            dim=dim,
+        ),
+    ]
+
+    schema = CollectionSchema(fields=fields, description="reverse image search")
+    collection = Collection(name=collection_name, schema=schema)
+
+    index_params = {"metric_type": "L2", "index_type": "IVF_FLAT", "params": {"nlist": 2048}}
+    collection.create_index(field_name="embedding", index_params=index_params)
+
+    return collection

fis/utils/config.py

@@ -0,0 +1,16 @@
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+# Read environment variables from .env file.
+load_dotenv()
+
+DIR_ROOT = Path(__file__).resolve().parent.parent.parent
+DIR_DATA = DIR_ROOT / "data"
+DIR_SCRAPING = DIR_DATA / "scraping"
+DIR_SCRAPING_IMAGES = DIR_SCRAPING / "images"
+
+S3_BUCKET = "fashion-img-search"
+S3_BUCKET_IMAGES = f"s3://{S3_BUCKET}/images/"
+
+FILE_SCRAPING_DATA = "items.jsonl"

fis/utils/constants.py

@@ -0,0 +1 @@
+ORGANISATION = "FSDL-Fashion"

fis/utils/data/download_fashionpedia.py

@@ -0,0 +1,56 @@
+import os
+import urllib.request
+import zipfile
+
+from fis.utils.config import DIR_DATA
+
+# Download from S3
+RAW_TRAIN_IMAGES = "https://s3.amazonaws.com/ifashionist-dataset/images/train2020.zip"
+RAW_VAL_IMAGES = "https://s3.amazonaws.com/ifashionist-dataset/images/val_test2020.zip"
+RAW_TRAIN_ANNOTATIONS = "https://s3.amazonaws.com/ifashionist-dataset/annotations/instances_attributes_train2020.json"
+RAW_VAL_ANNOTATIONS = "https://s3.amazonaws.com/ifashionist-dataset/annotations/instances_attributes_val2020.json"
+
+# to local disk
+TRAIN_ANNOTATIONS = "train.json"
+VAL_ANNOTATIONS = "val.json"
+
+
+def download(url: str, target: str) -> None:
+    """Download image and annotations.
+
+    Args:
+        url: url to download from.
+        target: file or directory to download to.
+    """
+    print(f"Downloading from {url}")
+
+    # Images
+    if url.split(".")[-1] == "zip":
+        path, _ = urllib.request.urlretrieve(url=url)  # noqa
+        with zipfile.ZipFile(path, "r") as f:
+            f.extractall(target)
+
+        os.remove(path)
+
+    # Annotations
+    else:
+        urllib.request.urlretrieve(url=url, filename=target)  # noqa
+
+
+def download_fashionpedia(destination_dir: str = DIR_DATA) -> None:
+    """Download the Fashionpedia dataset.
+
+    Args:
+        destination_dir: directory where the dataset will be saved.
+    """
+    os.makedirs(destination_dir, exist_ok=True)
+
+    download(url=RAW_TRAIN_ANNOTATIONS, target=os.path.join(destination_dir, TRAIN_ANNOTATIONS))
+    download(url=RAW_VAL_ANNOTATIONS, target=os.path.join(destination_dir, VAL_ANNOTATIONS))
+
+    download(url=RAW_TRAIN_IMAGES, target=destination_dir)
+    download(url=RAW_VAL_IMAGES, target=destination_dir)
+
+
+if __name__ == "__main__":
+    download_fashionpedia()

fis/utils/s3.py

@@ -0,0 +1,38 @@
+import io
+from typing import List
+
+import boto3
+from PIL import Image
+
+from fis.utils import config as cfg
+
+s3 = boto3.resource("s3")
+
+
+def list_images_from_bucket(bucket: str = cfg.S3_BUCKET) -> List[str]:
+    """List jpeg images from a bucket.
+
+    Args:
+        bucket: Name of the bucket. Defaults to cfg.S3_BUCKET.
+
+    Returns:
+        List of image names.
+    """
+    my_bucket = s3.Bucket(bucket)
+
+    images = []
+    for _object in my_bucket.objects.all():
+        key = _object.key
+        if ".jpg" in key:
+            images.append(key)
+
+    return images
+
+
+def read_image_from_s3(key, bucket: str = cfg.S3_BUCKET):
+
+    bucket = s3.Bucket(bucket)
+    image = bucket.Object(key)
+    img_data = image.get().get("Body").read()
+
+    return Image.open(io.BytesIO(img_data))