move database calls to flask app

app.py

@@ -1,15 +1,10 @@
 import gradio as gr
-import os
 from transformers import AutoImageProcessor, AutoModel
 import torch
-from pymongo import MongoClient
 from PIL import Image
 import json
 import numpy as np
 import faiss
-from dotenv import load_dotenv
-
-load_dotenv()
 
 
 # Init similarity search AI model and processor
@@ -17,11 +12,6 @@ torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 dino_v2_model = AutoModel.from_pretrained("./dinov2-base").to(torch_device)
 dino_v2_image_processor = AutoImageProcessor.from_pretrained("./dinov2-base")
 
-# MongoDB
-MONGO_URI = os.environ.get("MONGO_URI")
-mongo = MongoClient(MONGO_URI)
-db = mongo["xbgp"]
-
 
 def process_image(image):
     """
@@ -68,35 +58,17 @@ def process_image(image):
     # Read the index file and perform search of top 50 images
     index = faiss.read_index("vector.index")
     distances, indices = index.search(vector, 50)
-    matches = ""
+
+    matches = []
+
     for idx, matching_gamerpic in enumerate(indices[0]):
-        gamerpic = images[matching_gamerpic]
-        # Return the corresponding title with only the matched gamerpic
-        title = db.titles.find_one(
-            {"gamerpics.cdn": gamerpic},
-            {"name": 1, "type": 1, "url": 1, "gamerpics.$": 1},
-        )
-        print(title["name"])
-        title["rank"] = idx
-        title["score"] = str(round((1 / (distances[0][idx] + 1) * 100), 2)) + "%"
-
-        html = f"""
-        <h3 class="mr-4 inline align-middle text-3xl hover:underline">Matching gamerpics: Top 50 results</h3>
-        <div class="mt-8 flex flex-wrap gap-x-2">
-            <a href="{title['url']}" alt="{title['name']}" class="min-w-[130px] grow">
-                <div id="{title['_id']}" hx-swap="morph:innerHTML" class="mb-4 rounded-xl border border-black/5 p-2 hover:border-transparent hover:bg-black/5">
-                    <div class="flex-grow items-center text-center">
-                        <img src="https://assets.xboxgamer.pics{title['gamerpics'][0]['cdn']}" width="64" height="64" class="mx-auto" alt="Gamerpic" />
-                        <span class="text-center align-middle text-lg">{title["name"]}</span>
-                        <div class="inline-block rounded-2xl border border-stone-200 bg-white px-2 py-1 text-xs font-bold uppercase">
-                            Score: {title["score"]}
-                        </div>
-                    </div>
-                </div>
-            </a>
-        </div>
-        """
-        matches += html
+        gamerpic = {}
+        gamerpic["cdn"] = images[matching_gamerpic]
+        gamerpic["score"] = str(round((1 / (distances[0][idx] + 1) * 100), 2)) + "%"
+
+        print(gamerpic)
+
+        matches.append(gamerpic)
 
     return matches
 
@@ -105,7 +77,7 @@ def process_image(image):
 iface = gr.Interface(
     fn=process_image,
     inputs=gr.Image(type="pil"),  # Adjust the shape as needed
-    outputs="html",  # Or any other output format that suits your needs
+    outputs="json",  # Or any other output format that suits your needs
 )
 
 # Launch the Gradio app

requirements.txt

@@ -4,6 +4,4 @@ torch==2.1.1+cpu
 numpy==1.26.0
 pillow==10.0.1
 transformers==4.34.0
-pymongo[srv]==3.11
-faiss-cpu==1.7.4
-python-dotenv
+faiss-cpu==1.7.4