Update main.py

main.py

@@ -1,21 +1,12 @@
 import os
 import json
 import base64
-import httpx
 import google.generativeai as genai
 
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel, HttpUrl
+from fastapi import FastAPI, HTTPException, File, UploadFile
 
 app = FastAPI(title="Business Card Analyzer")
-
-# Input model for the API
-class ImageRequest(BaseModel):
-    image_url: HttpUrl
-
-# Configure the Google Generative AI client using an environment variable.
-# Make sure to set the GOOGLE_API_KEY variable when running the app.
-GOOGLE_API_KEY = "AIzaSyDcYyq3w21iwipYn17wCAQo3AYWhUIGDSI"
+GOOGLE_API_KEY = 'AIzaSyDcYyq3w21iwipYn17wCAQo3AYWhUIGDSI'
 if not GOOGLE_API_KEY:
     raise Exception("Please set the GOOGLE_API_KEY environment variable.")
 
@@ -37,19 +28,26 @@ Analyze the business card and extract the following information in JSON format:
 Ensure that the response is strictly in JSON format and includes all the fields mentioned above. If any information is missing or unclear, set the corresponding field to null.
 """
 
-@app.post("/analyze", summary="Analyze a business card image")
-async def analyze_business_card(request: ImageRequest):
-    # Fetch the image from the provided URL
+@app.post("/analyze", summary="Analyze a business card image file")
+async def analyze_business_card(file: UploadFile = File(...)):
+    # Validate the file type. Only JPEG and PNG are supported.
+    if file.content_type not in ["image/jpeg", "image/png"]:
+        raise HTTPException(
+            status_code=400,
+            detail="Invalid file type. Only JPEG and PNG images are supported."
+        )
+    
+    # Read the uploaded file
     try:
-        async with httpx.AsyncClient() as client:
-            response = await client.get(str(request.image_url))
-            response.raise_for_status()
-            img_bytes = response.content
+        file_content = await file.read()
     except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Failed to fetch image: {e}")
-
+        raise HTTPException(
+            status_code=400,
+            detail=f"Failed to read the uploaded file: {e}"
+        )
+    
     # Encode the image in base64
-    encoded_image = base64.b64encode(img_bytes).decode("utf-8")
+    encoded_image = base64.b64encode(file_content).decode("utf-8")
 
     # Create an instance of the Gemini model
     model = genai.GenerativeModel(model_name="gemini-1.5-pro")
@@ -58,7 +56,7 @@ async def analyze_business_card(request: ImageRequest):
     try:
         result = model.generate_content([
             {
-                "mime_type": "image/jpeg",
+                "mime_type": file.content_type,
                 "data": encoded_image,
             },
             PROMPT,
@@ -66,7 +64,7 @@ async def analyze_business_card(request: ImageRequest):
         # Get the raw text response
         raw = result.text.strip()
 
-        # Remove markdown formatting if present (i.e. ```json ... ```)
+        # Remove markdown formatting if present (e.g., ```json ... ```)
         if raw.startswith("```"):
             lines = raw.splitlines()
             if lines[0].startswith("```"):
@@ -77,9 +75,15 @@ async def analyze_business_card(request: ImageRequest):
 
         data = json.loads(raw)
     except json.JSONDecodeError as jde:
-        raise HTTPException(status_code=500, detail=f"JSON decoding error: {jde}. Response text was: {raw}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"JSON decoding error: {jde}. Response text was: {raw}"
+        )
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"An error occurred while processing the image: {e}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"An error occurred while processing the image: {e}"
+        )
 
     return data