uploaded baseline app

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+embeddings.index filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,95 @@
+import os
+import boto3
+import streamlit as st
+import faiss
+import pandas as pd
+from PIL import Image
+from model_prediction import Ranker
+from io import BytesIO
+
+
+@st.cache
+def load_model():
+    return Ranker()
+
+
+def load_faiss_index():
+    return faiss.read_index('embeddings.index')
+
+
+def load_labels():
+    return pd.read_csv("labels.csv")
+
+
+class ModelLoader:
+    model = None
+    index = None
+    labels = None
+
+    @classmethod
+    def get_model(cls):
+        if cls.model is None:
+            cls.model = load_model()
+        return cls.model
+
+    @classmethod
+    def get_index(cls):
+        if cls.index is None:
+            cls.index = load_faiss_index()
+        return cls.index
+
+    @classmethod
+    def get_labels(cls):
+        if cls.labels is None:
+            cls.labels = load_labels()
+        return cls.labels
+
+
+target_size = (224, 224)
+st.set_page_config(page_title="Product Retrieval App")
+st.title("Product Retrieval App")
+
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+loading_text = st.empty()
+
+s3 = boto3.client(
+    's3',
+    aws_access_key_id='AKIAUUWYLZEQYT6ESW4Q',
+    aws_secret_access_key='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',
+    region_name='eu-west-1'
+)
+
+bucket_name = "product-retrieval"
+
+if uploaded_file is not None:
+    image = Image.open(uploaded_file)
+    st.image(image, caption="Uploaded image", use_column_width=True)
+
+    loading_text.text("Loading predictions...")
+
+    model = ModelLoader.get_model()
+    index = ModelLoader.get_index()
+    labels = ModelLoader.get_labels()
+
+    image_embedding = model.predict(image)
+    distances, indices = index.search(image_embedding, 12)
+    predicted_images = labels["path"][indices[0]].to_list()
+    loading_text.empty()
+
+    col1, col2, col3, col4 = st.columns(4)
+
+    for i, img_path in enumerate(predicted_images):
+        response = s3.get_object(Bucket=bucket_name, Key=img_path.split("/")[-1])
+        image_data = response['Body'].read()
+        img = Image.open(BytesIO(image_data)).resize(target_size)
+
+        if i % 4 == 0:
+            column = col1
+        elif i % 4 == 1:
+            column = col2
+        elif i % 4 == 2:
+            column = col3
+        else:
+            column = col4
+        with column:
+            st.image(img, caption=f"Predicted image {i+1}", use_column_width=True)

model_prediction.py

@@ -0,0 +1,88 @@
+from sklearn.preprocessing import normalize
+import torchvision.transforms as T
+import open_clip
+import torch
+import math
+from torch import nn
+import torch.nn.functional as F
+
+
+def get_final_transform():
+    final_transform = T.Compose([
+        T.Resize(
+            size=(224, 224),
+            interpolation=T.InterpolationMode.BICUBIC,
+            antialias=True),
+        T.ToTensor(),
+        T.Normalize(
+            mean=(0.48145466, 0.4578275, 0.40821073),
+            std=(0.26862954, 0.26130258, 0.27577711)
+        )
+    ])
+    return final_transform
+
+
+class Clip_Products(nn.Module):
+    def __init__(self, vit_backbone, head_size, k=3):
+        super(Clip_Products, self).__init__()
+        self.head = HeadV2(head_size, k)
+        self.encoder = vit_backbone.visual
+
+    def forward(self, x):
+        x = self.encoder(x)
+        return self.head(x)
+
+
+class ArcMarginProduct_subcenter(nn.Module):
+    def __init__(self, in_features, out_features, k=3):
+        super().__init__()
+        self.weight = nn.Parameter(torch.FloatTensor(out_features * k, in_features))
+        self.reset_parameters()
+        self.k = k
+        self.out_features = out_features
+
+    def reset_parameters(self):
+        stdv = 1. / math.sqrt(self.weight.size(1))
+        self.weight.data.uniform_(-stdv, stdv)
+
+    def forward(self, features):
+        cosine_all = F.linear(F.normalize(features), F.normalize(self.weight))
+        cosine_all = cosine_all.view(-1, self.out_features, self.k)
+        cosine, _ = torch.max(cosine_all, dim=2)
+        return cosine
+
+
+class HeadV2(nn.Module):
+    def __init__(self, hidden_size, k=3):
+        super(HeadV2, self).__init__()
+        self.arc = ArcMarginProduct_subcenter(hidden_size, 9691, k)
+
+    def forward(self, x):
+        output = self.arc(x)
+        return output, F.normalize(x)
+
+
+class Ranker:
+    def __init__(self):
+        self.model_path = "model/best_model.pt"
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+        backbone, _, _ = open_clip.create_model_and_transforms('ViT-L-14', None)
+        self.model = Clip_Products(backbone, 768, 3)
+
+        checkpoint = torch.load(self.model_path, map_location=self.device)
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        self.model.to(self.device)
+
+    def predict(self, img):
+        transform_img = get_final_transform()
+        query = transform_img(img)
+
+        with torch.no_grad():
+            self.model.eval()
+
+            images = query.to(self.device, dtype=torch.float).unsqueeze(0)
+            _, embeddings = self.model(images)
+
+        query_embeddings = embeddings.detach().cpu().numpy()
+        return normalize(query_embeddings)

requirements.txt

@@ -0,0 +1,8 @@
+boto3
+faiss-cpu
+pandas
+Pillow
+scikit-learn
+torchvision
+torch
+open_clip_torch