Upload predict.py

img2art_search/models/predict.py

@@ -11,13 +11,10 @@ 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)
+        img = img.convert("RGB")
+        pred_image_data = transform(img).unsqueeze(0).to(DEVICE)
         indices, distances = search_image(pred_image_data)
         results = []
         for index, distance in zip(indices, distances):
@@ -31,7 +28,6 @@ def predict(img: Image.Image) -> list:
                     str(distance),
                 )
             )
-        os.remove(tmp_img_path)
         return results
     else:
         return []