Update app.py

app.py

@@ -6,13 +6,13 @@ import base64
 import json
 import uvicorn
 from typing import Dict
+from PIL import Image
+import io
 
-# Configure Google API
 GOOGLE_API_KEY = 'AIzaSyDcYyq3w21iwipYn17wCAQo3AYWhUIGDSI'
 genai.configure(api_key=GOOGLE_API_KEY)
 model = genai.GenerativeModel(model_name="gemini-1.5-pro")
 
-# Initialize FastAPI
 app = FastAPI()
 app.add_middleware(
     CORSMiddleware,
@@ -21,9 +21,16 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-def process_image(image_data: bytes) -> Dict:
-    # Encode image to base64
-    encoded_image = base64.b64encode(image_data).decode("utf-8")
+def process_image(image) -> Dict:
+    if isinstance(image, str):
+        image = Image.open(image)
+    
+    # Convert PIL Image to bytes
+    img_byte_arr = io.BytesIO()
+    image.save(img_byte_arr, format='JPEG')
+    img_byte_arr = img_byte_arr.getvalue()
+    
+    encoded_image = base64.b64encode(img_byte_arr).decode("utf-8")
     
     prompt = """
     Analyze the business card and extract the following information in JSON format:
@@ -44,7 +51,6 @@ def process_image(image_data: bytes) -> Dict:
         prompt,
     ])
     
-    # Clean response and parse JSON
     raw = response.text.strip()
     if raw.startswith("```"):
         lines = raw.splitlines()
@@ -56,14 +62,13 @@ def process_image(image_data: bytes) -> Dict:
     
     return json.loads(raw)
 
-# FastAPI endpoint
 @app.post("/extract")
 async def extract_info(file: UploadFile = File(...)):
     contents = await file.read()
-    result = process_image(contents)
+    img = Image.open(io.BytesIO(contents))
+    result = process_image(img)
     return result
 
-# Gradio interface
 def gradio_process(image):
     if image is None:
         return json.dumps({"error": "No image provided"}, indent=2)
@@ -72,13 +77,12 @@ def gradio_process(image):
 
 interface = gr.Interface(
     fn=gradio_process,
-    inputs=gr.Image(type="binary", label="Upload Business Card"),
+    inputs=gr.Image(type="pil", label="Upload Business Card"),
     outputs=gr.JSON(label="Extracted Information"),
     title="Business Card Information Extractor",
     description="Upload a business card image to extract contact information"
 )
 
-# Mount Gradio app to FastAPI
 app = gr.mount_gradio_app(app, interface, path="/")
 
 if __name__ == "__main__":