byaldi

app.py

@@ -1,6 +1,9 @@
+import base64
+import io
 import json
 import os
 from datetime import datetime
+from io import BytesIO
 
 import dotenv
 import lancedb
@@ -8,6 +11,7 @@ import requests
 from datasets import load_dataset
 from fasthtml.common import *  # noqa
 from huggingface_hub import login, whoami
+from PIL import Image
 from rerankers import Reranker
 
 
@@ -179,7 +183,7 @@ style = Style("""
             margin-top: 20px;
         }
         .image-result {
-            width: calc(20% - 10px);
+            width: calc(33% - 10px);
             text-align: center;
         }
         .image-result img {
@@ -253,9 +257,38 @@ def SearchResult(result):
     )
 
 
+def base64_to_pil(base64_string):
+    # Remove the "data:image/png;base64," part if it exists
+    if "base64," in base64_string:
+        base64_string = base64_string.split("base64,")[1]
+
+    # Decode the base64 string
+    img_data = base64.b64decode(base64_string)
+
+    # Open the image using PIL
+    img = Image.open(BytesIO(img_data))
+
+    return img
+
+
+def process_image(image, max_size=(500, 500), quality=85):
+    pil_image = base64_to_pil(image)
+    img_byte_arr = io.BytesIO()
+    pil_image.thumbnail(max_size)
+    pil_image.save(img_byte_arr, format="JPEG", quality=quality, optimize=True)
+    return f"data:image/jpeg;base64,{base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')}"
+
+
+# def ImageResult(image):
+#     return Div(
+#         Img(src=f"data:image/jpeg;base64,{image}", alt="arxiv image"),
+#         cls="image-result",
+#     )
+
+
 def ImageResult(image):
     return Div(
-        Img(src=f"data:image/jpeg;base64,{image}", alt="arxiv image"),
+        Img(src=process_image(image), alt="arxiv image"),
         cls="image-result",
     )
 
@@ -277,8 +310,10 @@ async def post(query: str):
     log_query_and_results(query, results)
 
     return Div(
+        Br(),
         H3("Byaldi Results"),
         Div(*[ImageResult(img) for img in image_results], cls="image-results"),
+        Br(),
         H3("Text Results"),
         Div(*[SearchResult(r) for r in results], id="text-results"),
         id="search-results",