from fastapi import FastAPI, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
import io
from transformers import pipeline

app = FastAPI()

# Enable CORS
origins = ["*"]  # Set this to your frontend URLs in production
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
    allow_headers=["*"],
)

# Define your VQA pipeline
vqa_pipeline = pipeline("visual-question-answering", model="Salesforce/blip-vqa-capfilt-large")


@app.post("/answer_question")
async def answer_question(question: str = Form(...), image: UploadFile = File(...)):
    # Read image file and convert to RGB format
    image_content = await image.read()
    image = Image.open(io.BytesIO(image_content)).convert("RGB")

    # Call the VQA pipeline to answer the question
    results = vqa_pipeline(question=question, image=image)

    # Assuming result is a list, get the first item
    if results:
        answer = results[0].get("answer", "No answer found")
    else:
        answer = "No answer found"

    return {"answer": answer}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)