Update app.py

app.py

@@ -1,25 +1,26 @@
 
-from fastapi import FastAPI, UploadFile
+from fastapi import FastAPI, File, UploadFile
 import pdfplumber
 import docx
 import openpyxl
 from pptx import Presentation
+import easyocr
 from transformers import pipeline
 import gradio as gr
 from fastapi.responses import RedirectResponse
 
-# ✅ Initialize FastAPI
+# Initialize FastAPI
 app = FastAPI()
 
-# ✅ Load AI Model
-qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-large", tokenizer="google/flan-t5-large", use_fast=True, device=0)
+# Load AI Model for Question Answering
+qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-large", tokenizer="google/flan-t5-large", use_fast=True)
 
-# ✅ Function to truncate text to avoid token limit errors
+# Function to truncate text to 450 tokens
 def truncate_text(text, max_tokens=450):
     words = text.split()
     return " ".join(words[:max_tokens])
 
-# ✅ Functions to extract text from different file formats
+# Functions to extract text from different file formats
 def extract_text_from_pdf(pdf_file):
     text = ""
     with pdfplumber.open(pdf_file) as pdf:
@@ -48,7 +49,12 @@ def extract_text_from_excel(excel_file):
             text.append(" ".join(map(str, row)))
     return "\n".join(text)
 
-# ✅ Function to answer questions based on document content
+def extract_text_from_image(image_file):
+    reader = easyocr.Reader(["en"])
+    result = reader.readtext(image_file)
+    return " ".join([res[1] for res in result])
+
+# Function to answer questions based on document content
 def answer_question_from_document(file, question):
     file_ext = file.name.split(".")[-1].lower()
     
@@ -65,13 +71,26 @@ def answer_question_from_document(file, question):
 
     if not text:
         return "No text extracted from the document."
+    
+    truncated_text = truncate_text(text)
+    input_text = f"Question: {question} Context: {truncated_text}"
+    response = qa_pipeline(input_text)
+    
+    return response[0]["generated_text"]
 
-    truncated_text = truncate_text(text)  # ✅ Prevents token limit error
+# Function to answer questions based on image content
+def answer_question_from_image(image, question):
+    image_text = extract_text_from_image(image)
+    if not image_text:
+        return "No text detected in the image."
+    
+    truncated_text = truncate_text(image_text)
     input_text = f"Question: {question} Context: {truncated_text}"
     response = qa_pipeline(input_text)
+    
     return response[0]["generated_text"]
 
-# ✅ Gradio UI for Document QA
+# Gradio UI for Document & Image QA
 doc_interface = gr.Interface(
     fn=answer_question_from_document,
     inputs=[gr.File(label="Upload Document"), gr.Textbox(label="Ask a Question")],
@@ -79,8 +98,15 @@ doc_interface = gr.Interface(
     title="AI Document Question Answering"
 )
 
-# ✅ Mount Gradio Interfaces
-demo = gr.TabbedInterface([doc_interface], ["Document QA"])
+img_interface = gr.Interface(
+    fn=answer_question_from_image,
+    inputs=[gr.Image(label="Upload Image"), gr.Textbox(label="Ask a Question")],
+    outputs="text",
+    title="AI Image Question Answering"
+)
+
+# Mount Gradio Interfaces
+demo = gr.TabbedInterface([doc_interface, img_interface], ["Document QA", "Image QA"])
 app = gr.mount_gradio_app(app, demo, path="/")
 
 @app.get("/")