Update app.py

app.py

@@ -3,7 +3,7 @@ import os
 import sys
 
 from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import RedirectResponse, JSONResponse
+from fastapi.responses import RedirectResponse, FileResponse
 import gradio as gr
 import requests
 import uvicorn
@@ -14,6 +14,8 @@ from PIL import Image
 from torch.utils.data import DataLoader
 from transformers import AutoProcessor
 import base64
+from reportlab.pdfgen import canvas
+from reportlab.lib.pagesizes import letter
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), './colpali-main')))
 
@@ -87,13 +89,25 @@ async def search(query: str, k: int):
 
     results = []
     for idx in top_k_indices:
-        img_byte_arr = io.BytesIO()
-        images[idx].save(img_byte_arr, format='PNG')
+        results.append({"image": images[idx], "page": f"Page {idx}"})
+
+    # Generate PDF
+    pdf_buffer = BytesIO()
+    c = canvas.Canvas(pdf_buffer, pagesize=letter)
+    width, height = letter
+    for result in results:
+        img = result["image"]
+        img_byte_arr = BytesIO()
+        img.save(img_byte_arr, format='PNG')
         img_byte_arr = img_byte_arr.getvalue()
         img_base64 = base64.b64encode(img_byte_arr).decode('utf-8')
-        results.append({"image": img_base64, "page": f"Page {idx}"})
+        c.drawImage(BytesIO(base64.b64decode(img_base64)), 0, 0, width, height)
+        c.showPage()
+
+    c.save()
+    pdf_buffer.seek(0)
 
-    return JSONResponse(content={"results": results})
+    return FileResponse(pdf_buffer, media_type='application/pdf', filename='search_results.pdf')
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)