updated code to recalculate embeddings

Devasy23 · Jul 30, 2024 · 02bd51e · 02bd51e
1 parent 1b7c685
commit 02bd51e
Showing 1 changed file with 4 additions and 7 deletions.
diff --git a/API/route.py b/API/route.py
@@ -38,7 +38,7 @@
 
 collection = 'faceEntries'
 collection2 = 'ImageDB'
-
+collection3 = 'VectorDB'
 
 # Models  for the data to be sent and received by the server
 class Employee(BaseModel):
@@ -112,14 +112,11 @@ async def recalculate_embeddings():
         if 'Images' in employee:
             images = employee['Images']
         else:
-            images = employee['Image']
+            images = [employee['Image']]
 
         for encoded_image in images:
-            encoded_image = encoded_image.replace('data:image/png;base64,', '')
-            encoded_image = encoded_image.strip()
-            encoded_image += '=' * (-len(encoded_image) % 4)
-            img_recovered = base64.b64decode(encoded_image)
-            pil_image = Image.open(BytesIO(img_recovered))
+
+            pil_image = Image.open(BytesIO(base64.b64decode(encoded_image)))
             image_filename = f'{employee["Name"]}.png'
             pil_image.save(image_filename)
             logging.debug(f'Image saved {employee["Name"]}')