VISION / main.py
arevedudaa's picture
Update main.py
23d6896 verified
import os
import json
import base64
import google.generativeai as genai
from fastapi import FastAPI, HTTPException, File, UploadFile
app = FastAPI(title="Business Card Analyzer")
GOOGLE_API_KEY = 'AIzaSyDcYyq3w21iwipYn17wCAQo3AYWhUIGDSI'
if not GOOGLE_API_KEY:
raise Exception("Please set the GOOGLE_API_KEY environment variable.")
genai.configure(api_key=GOOGLE_API_KEY)
# Define the prompt to instruct the model.
PROMPT = """
Analyze the business card and extract the following information in JSON format:
{
"name": "Full Name",
"mobile_number": "Mobile Number",
"company_name": "Company Name",
"email": "Email Address",
"job_title": "Job Title",
"address": "Address",
"website": "Website URL"
}
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 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:
file_content = await file.read()
except Exception as e:
raise HTTPException(
status_code=400,
detail=f"Failed to read the uploaded file: {e}"
)
# Encode the image in base64
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")
# Call the model and process its response
try:
result = model.generate_content([
{
"mime_type": file.content_type,
"data": encoded_image,
},
PROMPT,
])
# Get the raw text response
raw = result.text.strip()
# Remove markdown formatting if present (e.g., ```json ... ```)
if raw.startswith("```"):
lines = raw.splitlines()
if lines[0].startswith("```"):
lines = lines[1:]
if lines and lines[-1].startswith("```"):
lines = lines[:-1]
raw = "\n".join(lines).strip()
data = json.loads(raw)
except json.JSONDecodeError as jde:
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}"
)
return data
if __name__ == "__main__":
import uvicorn
uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)