main.py

import fitz
from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import JSONResponse
from transformers import pipeline
from PIL import Image
from io import BytesIO
from starlette.middleware import Middleware
from starlette.middleware.cors import CORSMiddleware
from pdf2image import convert_from_bytes

app = FastAPI()

# Set up CORS middleware
origins = ["*"]  # or specify your list of allowed origins
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

nlp_qa = pipeline("document-question-answering", model="jinhybr/OCR-DocVQA-Donut")
nlp_qa_v2 = pipeline("document-question-answering", model="faisalraza/layoutlm-invoices")
nlp_qa_v3 = pipeline("question-answering", model="deepset/roberta-base-squad2")
nlp_classification = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")

description = """
## Image-based Document QA
This API performs document question answering using a LayoutLMv2-based model.

### Endpoints:
- **POST /uploadfile/:** Upload an image file to extract text and answer provided questions.
- **POST /pdfQA/:** Provide a PDF file to extract text and answer provided questions.
"""

app = FastAPI(docs_url="/", description=description)

@app.post("/uploadfile/", description="Upload an image file to extract text and answer provided questions.")
async def perform_document_qa(
    file: UploadFile = File(...),
    questions: str = Form(...),
):
    try:
        # Read the uploaded file as bytes
        contents = await file.read()

        # Open the image using PIL
        image = Image.open(BytesIO(contents))

        # Perform document question answering for each question using LayoutLMv2-based model
        answers_dict = {}
        for question in questions.split(','):
            result = nlp_qa(
                image,
                question.strip()
            )

            # Access the 'answer' key from the first item in the result list
            answer = result[0]['answer']

            # Format the question as a string without extra characters
            formatted_question = question.strip("[]")

            answers_dict[formatted_question] = answer

        return answers_dict
    except Exception as e:
        return JSONResponse(content=f"Error processing file: {str(e)}", status_code=500)

@app.post("/uploadfilev2/", description="Upload an image file to extract text and answer provided questions.")
async def perform_document_qa(
    file: UploadFile = File(...),
    questions: str = Form(...),
):
    try:
        # Read the uploaded file as bytes
        contents = await file.read()

        # Open the image using PIL
        image = Image.open(BytesIO(contents))

        # Perform document question answering for each question using LayoutLMv2-based model
        answers_dict = {}
        for question in questions.split(','):
            result = nlp_qa_v2(
                image,
                question.strip()
            )

            # Access the 'answer' key from the first item in the result list
            answer = result[0]['answer']

            # Format the question as a string without extra characters
            formatted_question = question.strip("[]")

            answers_dict[formatted_question] = answer

        return answers_dict
    except Exception as e:
        return JSONResponse(content=f"Error processing file: {str(e)}", status_code=500)

@app.post("/uploadfilev3/", description="Upload an image file to extract text and answer provided questions.")
async def perform_document_qa(
    context: str = Form(...),
    question: str = Form(...),
):
    try:
        QA_input = {
            'question': question,
            'context': context
        }
        
        res = nlp_qa_v3(QA_input)

        return res['answer']
    except Exception as e:
        return JSONResponse(content=f"Error processing file: {str(e)}", status_code=500)

@app.post("/classify/", description="Classify the provided text.")
async def classify_text(text: str = Form(...)):
    try:
        # Perform text classification using the pipeline
        result = nlp_classification(text)

        # Return the classification result
        return result
    except Exception as e:
        return JSONResponse(content=f"Error classifying text: {str(e)}", status_code=500)


# Set up CORS middleware
origins = ["*"]  # or specify your list of allowed origins
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)