Update app.py

app.py

@@ -69,125 +69,128 @@ app = gr.mount_gradio_app(app, demo, path="/")
 def root():
     return RedirectResponse(url="/")
 """
-from fastapi import FastAPI
-from fastapi.responses import RedirectResponse
 import gradio as gr
-import pytesseract
-from PIL import Image
-import fitz  # PyMuPDF
-import pdfplumber
-import easyocr
-import docx
-import openpyxl
-from pptx import Presentation
-from transformers import pipeline
+import fitz  # PyMuPDF for PDFs
+import easyocr  # OCR for images
+import openpyxl  # XLSX processing
+import pptx  # PPTX processing
+import docx  # DOCX processing
+import json  # Exporting results
 from deep_translator import GoogleTranslator
-import json
-import os
+from transformers import pipeline
+from fastapi import FastAPI
+from starlette.responses import RedirectResponse
 
+# Initialize FastAPI app
 app = FastAPI()
 
-qa_pipeline = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")
-reader = easyocr.Reader(['en'])
-
-# Utility functions
+# Initialize AI Models
+qa_model = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
+image_captioning = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
+reader = easyocr.Reader(['en', 'fr'])  # EasyOCR for image text extraction (English & French)
 
+# ---- TEXT EXTRACTION FUNCTIONS ----
 def extract_text_from_pdf(pdf_file):
-    with pdfplumber.open(pdf_file) as pdf:
-        return "\n".join([page.extract_text() for page in pdf.pages if page.extract_text()])
+    """Extract text from a PDF file."""
+    text = []
+    try:
+        with fitz.open(pdf_file) as doc:
+            for page in doc:
+                text.append(page.get_text("text"))
+    except Exception as e:
+        return f"Error reading PDF: {e}"
+    return "\n".join(text)
 
 def extract_text_from_docx(docx_file):
+    """Extract text from a DOCX file."""
     doc = docx.Document(docx_file)
-    return "\n".join([para.text for para in doc.paragraphs])
+    return "\n".join([p.text for p in doc.paragraphs if p.text.strip()])
 
 def extract_text_from_pptx(pptx_file):
-    prs = Presentation(pptx_file)
+    """Extract text from a PPTX file."""
     text = []
-    for slide in prs.slides:
-        for shape in slide.shapes:
-            if hasattr(shape, "text"):
-                text.append(shape.text)
+    try:
+        presentation = pptx.Presentation(pptx_file)
+        for slide in presentation.slides:
+            for shape in slide.shapes:
+                if hasattr(shape, "text"):
+                    text.append(shape.text)
+    except Exception as e:
+        return f"Error reading PPTX: {e}"
     return "\n".join(text)
 
 def extract_text_from_xlsx(xlsx_file):
-    wb = openpyxl.load_workbook(xlsx_file)
+    """Extract text from an XLSX file."""
     text = []
-    for sheet in wb.worksheets:
-        for row in sheet.iter_rows():
-            text.extend([str(cell.value) for cell in row if cell.value is not None])
+    try:
+        wb = openpyxl.load_workbook(xlsx_file)
+        for sheet in wb.sheetnames:
+            ws = wb[sheet]
+            for row in ws.iter_rows(values_only=True):
+                text.append(" ".join(str(cell) for cell in row if cell))
+    except Exception as e:
+        return f"Error reading XLSX: {e}"
     return "\n".join(text)
 
-def extract_text(file):
-    ext = os.path.splitext(file.name)[1].lower()
-    if ext == ".pdf":
-        return extract_text_from_pdf(file)
-    elif ext == ".docx":
-        return extract_text_from_docx(file)
-    elif ext == ".pptx":
-        return extract_text_from_pptx(file)
-    elif ext == ".xlsx":
-        return extract_text_from_xlsx(file)
-    else:
-        return "Unsupported file type"
-
-def answer_question_from_doc(file, question, translate_to="en"):
-    context = extract_text(file)
-    result = qa_pipeline(question=question, context=context)
-    translated = GoogleTranslator(source='auto', target=translate_to).translate(result["answer"])
-    return {
-        "answer": translated,
-        "score": result["score"],
-        "original": result["answer"]
-    }
-
-def answer_question_from_image(image, question, translate_to="en"):
-    img_text = pytesseract.image_to_string(image)
-    if not img_text.strip():
-        img_text = "\n".join([line[1] for line in reader.readtext(image)])
-    result = qa_pipeline(question=question, context=img_text)
-    translated = GoogleTranslator(source='auto', target=translate_to).translate(result["answer"])
-    return {
-        "answer": translated,
-        "score": result["score"],
-        "original": result["answer"]
-    }
+def extract_text_from_image(image_path):
+    """Extract text from an image using EasyOCR."""
+    result = reader.readtext(image_path, detail=0)
+    return " ".join(result)  # Return text as a single string
 
-# Gradio Interfaces
+# ---- MAIN PROCESSING FUNCTIONS ----
+def answer_question_from_doc(file, question):
+    """Process document and answer a question based on its content."""
+    ext = file.name.split(".")[-1].lower()
+    
+    if ext == "pdf":
+        context = extract_text_from_pdf(file.name)
+    elif ext == "docx":
+        context = extract_text_from_docx(file.name)
+    elif ext == "pptx":
+        context = extract_text_from_pptx(file.name)
+    elif ext == "xlsx":
+        context = extract_text_from_xlsx(file.name)
+    else:
+        return "Unsupported file format."
 
-doc_interface = gr.Interface(
-    fn=answer_question_from_doc,
-    inputs=[
-        gr.File(label="Upload Document (PDF, DOCX, PPTX, XLSX)"),
-        gr.Textbox(label="Ask a Question"),
-        gr.Textbox(label="Translate Answer To (e.g., en, fr, ar)", value="en")
-    ],
-    outputs=[
-        gr.Textbox(label="Translated Answer"),
-        gr.Number(label="Confidence Score"),
-        gr.Textbox(label="Original Answer")
-    ],
-    title="📄 Document QA + Translation + Export"
-)
+    if not context.strip():
+        return "No text found in the document."
 
-img_interface = gr.Interface(
-    fn=answer_question_from_image,
-    inputs=[
-        gr.Image(label="Upload Image"),
-        gr.Textbox(label="Ask a Question"),
-        gr.Textbox(label="Translate Answer To (e.g., en, fr, ar)", value="en")
-    ],
-    outputs=[
-        gr.Textbox(label="Translated Answer"),
-        gr.Number(label="Confidence Score"),
-        gr.Textbox(label="Original Answer")
-    ],
-    title="🖼️ Image QA + OCR + Translation + Export"
-)
+    # Generate answer using AI
+    answer = qa_model(question + " " + context, max_length=100)[0]["generated_text"]
+    return answer
 
-# Combine interfaces
+def answer_question_from_image(image, question):
+    """Process an image, extract text, and answer a question."""
+    img_text = extract_text_from_image(image)
+    if not img_text.strip():
+        return "No readable text found in the image."
+
+    # Generate answer using AI
+    answer = qa_model(question + " " + img_text, max_length=100)[0]["generated_text"]
+    return answer
+
+# ---- GRADIO INTERFACES ----
+with gr.Blocks() as doc_interface:
+    gr.Markdown("## Document Question Answering")
+    file_input = gr.File(label="Upload DOCX, PPTX, XLSX, or PDF")
+    question_input = gr.Textbox(label="Ask a question")
+    answer_output = gr.Textbox(label="Answer")
+    file_submit = gr.Button("Get Answer")
+    file_submit.click(answer_question_from_doc, inputs=[file_input, question_input], outputs=answer_output)
+
+with gr.Blocks() as img_interface:
+    gr.Markdown("## Image Question Answering")
+    image_input = gr.Image(label="Upload an Image")
+    img_question_input = gr.Textbox(label="Ask a question")
+    img_answer_output = gr.Textbox(label="Answer")
+    image_submit = gr.Button("Get Answer")
+    image_submit.click(answer_question_from_image, inputs=[image_input, img_question_input], outputs=img_answer_output)
+
+# ---- MOUNT GRADIO APP ----
 demo = gr.TabbedInterface([doc_interface, img_interface], ["Document QA", "Image QA"])
 app = gr.mount_gradio_app(app, demo, path="/")
 
 @app.get("/")
-def root():
+def home():
     return RedirectResponse(url="/")