Update main.py

main.py

@@ -1,6 +1,8 @@
 from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.responses import JSONResponse
 from transformers import pipeline
+from PIL import Image
+from io import BytesIO
 
 app = FastAPI()
 
@@ -27,11 +29,14 @@ async def perform_document_qa(
         # 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 LayoutLM-based model
         answers_dict = {}
         for question in questions.split(','):
             result = nlp_qa(
-                contents,
+                image,
                 question.strip()
             )
             answers_dict[question] = result['answer']
@@ -49,11 +54,14 @@ async def load_file(
         # 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 LayoutLM-based model
         answers_dict = {}
         for question in questions.split(','):
             result = nlp_qa(
-                contents,
+                image,
                 question.strip()
             )
             answers_dict[question] = result['answer']