correct paths to raw images

app.py

@@ -22,8 +22,6 @@ from typing import Callable, Dict, List, Tuple
 from PIL.Image import Image
 
 print(__file__)
-import fashion_aggregator.fashion_aggregator as fa
-
 
 os.environ["CUDA_VISIBLE_DEVICES"] = ""  # do not use GPU
 
@@ -41,21 +39,21 @@ EMBEDDINGS_FILE = os.path.join(EMBEDDINGS_DIR, "embeddings.pkl")
 RAW_PHOTOS_DIR = "artifacts/raw-photos"
 
 # Download image embeddings and raw photos
-wandb.login(key=os.getenv('wandb'))
-api = wandb.Api()
-artifact_embeddings = api.artifact("ryparmar/fashion-aggregator/unimoda-images:v1")
-artifact_embeddings.download(EMBEDDINGS_DIR)
-artifact_raw_photos = api.artifact("ryparmar/fashion-aggregator/unimoda-raw-images:v1")
-artifact_raw_photos.download("artifacts")
+# wandb.login(key="4b5a23a662b20fdd61f2aeb5032cf56fdce278a4")  # os.getenv('wandb')
+# api = wandb.Api()
+# artifact_embeddings = api.artifact("ryparmar/fashion-aggregator/unimoda-images:v1")
+# artifact_embeddings.download(EMBEDDINGS_DIR)
+# artifact_raw_photos = api.artifact("ryparmar/fashion-aggregator/unimoda-raw-images:v1")
+# artifact_raw_photos.download("artifacts")
 
-with zipfile.ZipFile("artifacts/unimoda.zip", 'r') as zip_ref:
-    zip_ref.extractall(RAW_PHOTOS_DIR)
+# with zipfile.ZipFile("artifacts/unimoda.zip", 'r') as zip_ref:
+#     zip_ref.extractall(RAW_PHOTOS_DIR)
 
 
 class TextEncoder:
     """Encodes the given text"""
 
-    def __init__(self, model_path='M-CLIP/XLM-Roberta-Large-Vit-B-32'):
+    def __init__(self, model_path="M-CLIP/XLM-Roberta-Large-Vit-B-32"):
         self.model = pt_multilingual_clip.MultilingualCLIP.from_pretrained(model_path)
         self.tokenizer = AutoTokenizer.from_pretrained(model_path)
 
@@ -69,11 +67,11 @@ class TextEncoder:
 class ImageEnoder:
     """Encodes the given image"""
 
-    def __init__(self, model_path='clip-ViT-B-32'):
+    def __init__(self, model_path="clip-ViT-B-32"):
         self.model = SentenceTransformer(model_path)
 
     @torch.no_grad()
-    def encode(self, image: Image.Image) -> torch.Tensor:
+    def encode(self, image: Image) -> torch.Tensor:
         """Predict/infer text embedding for a given query."""
         image_emb = self.model.encode([image], convert_to_tensor=True, show_progress_bar=False)
         return image_emb
@@ -81,24 +79,28 @@ class ImageEnoder:
 
 class Retriever:
     """Retrieves relevant images for a given text embedding."""
-    
+
     def __init__(self, image_embeddings_path=None):
         self.text_encoder = TextEncoder()
         self.image_encoder = ImageEnoder()
 
-        with open(image_embeddings_path, 'rb') as file:
-            self.image_names, self.image_embeddings = pickle.load(file)  
+        with open(image_embeddings_path, "rb") as file:
+            self.image_names, self.image_embeddings = pickle.load(file)
+            self.image_names = [
+                img_name.replace("fashion-aggregator/fashion_aggregator/data/photos/", "")
+                for img_name in self.image_names
+            ]
         print("Images:", len(self.image_names))
 
     @torch.no_grad()
-    def predict(self, text_query: str, k: int=10) -> List[Any]:
+    def predict(self, text_query: str, k: int = 10) -> List[Any]:
         """Return top-k relevant items for a given embedding"""
         query_emb = self.text_encoder.encode(text_query)
         relevant_images = util.semantic_search(query_emb, self.image_embeddings, top_k=k)[0]
         return relevant_images
 
     @torch.no_grad()
-    def search_images(self, text_query: str, k: int=6) -> Dict[str, List[Any]]:
+    def search_images(self, text_query: str, k: int = 6) -> Dict[str, List[Any]]:
         """Return top-k relevant images for a given embedding"""
         images = self.predict(text_query, k)
         paths_and_scores = {"path": [], "score": []}
@@ -155,7 +157,7 @@ class PredictorBackend:
             self.url = url
             self._predict = self._predict_from_endpoint
         else:
-            model = fa.Retriever()
+            model = Retriever(image_embeddings_path=EMBEDDINGS_FILE)
             self._predict = model.predict
             self._search_images = model.search_images