methodw commited on
Commit
3989d63
1 Parent(s): ded8264

move database calls to flask app

Browse files
Files changed (2) hide show
  1. app.py +11 -39
  2. requirements.txt +1 -3
app.py CHANGED
@@ -1,15 +1,10 @@
1
  import gradio as gr
2
- import os
3
  from transformers import AutoImageProcessor, AutoModel
4
  import torch
5
- from pymongo import MongoClient
6
  from PIL import Image
7
  import json
8
  import numpy as np
9
  import faiss
10
- from dotenv import load_dotenv
11
-
12
- load_dotenv()
13
 
14
 
15
  # Init similarity search AI model and processor
@@ -17,11 +12,6 @@ torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
17
  dino_v2_model = AutoModel.from_pretrained("./dinov2-base").to(torch_device)
18
  dino_v2_image_processor = AutoImageProcessor.from_pretrained("./dinov2-base")
19
 
20
- # MongoDB
21
- MONGO_URI = os.environ.get("MONGO_URI")
22
- mongo = MongoClient(MONGO_URI)
23
- db = mongo["xbgp"]
24
-
25
 
26
  def process_image(image):
27
  """
@@ -68,35 +58,17 @@ def process_image(image):
68
  # Read the index file and perform search of top 50 images
69
  index = faiss.read_index("vector.index")
70
  distances, indices = index.search(vector, 50)
71
- matches = ""
 
 
72
  for idx, matching_gamerpic in enumerate(indices[0]):
73
- gamerpic = images[matching_gamerpic]
74
- # Return the corresponding title with only the matched gamerpic
75
- title = db.titles.find_one(
76
- {"gamerpics.cdn": gamerpic},
77
- {"name": 1, "type": 1, "url": 1, "gamerpics.$": 1},
78
- )
79
- print(title["name"])
80
- title["rank"] = idx
81
- title["score"] = str(round((1 / (distances[0][idx] + 1) * 100), 2)) + "%"
82
-
83
- html = f"""
84
- <h3 class="mr-4 inline align-middle text-3xl hover:underline">Matching gamerpics: Top 50 results</h3>
85
- <div class="mt-8 flex flex-wrap gap-x-2">
86
- <a href="{title['url']}" alt="{title['name']}" class="min-w-[130px] grow">
87
- <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">
88
- <div class="flex-grow items-center text-center">
89
- <img src="https://assets.xboxgamer.pics{title['gamerpics'][0]['cdn']}" width="64" height="64" class="mx-auto" alt="Gamerpic" />
90
- <span class="text-center align-middle text-lg">{title["name"]}</span>
91
- <div class="inline-block rounded-2xl border border-stone-200 bg-white px-2 py-1 text-xs font-bold uppercase">
92
- Score: {title["score"]}
93
- </div>
94
- </div>
95
- </div>
96
- </a>
97
- </div>
98
- """
99
- matches += html
100
 
101
  return matches
102
 
@@ -105,7 +77,7 @@ def process_image(image):
105
  iface = gr.Interface(
106
  fn=process_image,
107
  inputs=gr.Image(type="pil"), # Adjust the shape as needed
108
- outputs="html", # Or any other output format that suits your needs
109
  )
110
 
111
  # Launch the Gradio app
 
1
  import gradio as gr
 
2
  from transformers import AutoImageProcessor, AutoModel
3
  import torch
 
4
  from PIL import Image
5
  import json
6
  import numpy as np
7
  import faiss
 
 
 
8
 
9
 
10
  # Init similarity search AI model and processor
 
12
  dino_v2_model = AutoModel.from_pretrained("./dinov2-base").to(torch_device)
13
  dino_v2_image_processor = AutoImageProcessor.from_pretrained("./dinov2-base")
14
 
 
 
 
 
 
15
 
16
  def process_image(image):
17
  """
 
58
  # Read the index file and perform search of top 50 images
59
  index = faiss.read_index("vector.index")
60
  distances, indices = index.search(vector, 50)
61
+
62
+ matches = []
63
+
64
  for idx, matching_gamerpic in enumerate(indices[0]):
65
+ gamerpic = {}
66
+ gamerpic["cdn"] = images[matching_gamerpic]
67
+ gamerpic["score"] = str(round((1 / (distances[0][idx] + 1) * 100), 2)) + "%"
68
+
69
+ print(gamerpic)
70
+
71
+ matches.append(gamerpic)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
72
 
73
  return matches
74
 
 
77
  iface = gr.Interface(
78
  fn=process_image,
79
  inputs=gr.Image(type="pil"), # Adjust the shape as needed
80
+ outputs="json", # Or any other output format that suits your needs
81
  )
82
 
83
  # Launch the Gradio app
requirements.txt CHANGED
@@ -4,6 +4,4 @@ torch==2.1.1+cpu
4
  numpy==1.26.0
5
  pillow==10.0.1
6
  transformers==4.34.0
7
- pymongo[srv]==3.11
8
- faiss-cpu==1.7.4
9
- python-dotenv
 
4
  numpy==1.26.0
5
  pillow==10.0.1
6
  transformers==4.34.0
7
+ faiss-cpu==1.7.4