Update app.py

app.py

@@ -122,22 +122,22 @@ if __name__ == "__main__":
 import gradio as gr
 import uvicorn
 import numpy as np
-import fitz  # PyMuPDF
+import pymupdf 
 import tika
 import torch
 from fastapi import FastAPI
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer
 from PIL import Image
 from io import BytesIO
 from starlette.responses import RedirectResponse
 from tika import parser
 from openpyxl import load_workbook
+from pptx import Presentation
+import easyocr
 import os
-import pymupdf
-# Initialize Tika for DOCX & PPTX parsing
-tika.initVM()
 
-# Initialize FastAPI
+tika.initVM() 
+
 app = FastAPI()
 
 # Load models
@@ -145,13 +145,15 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 qa_pipeline = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", device=device)
 image_captioning_pipeline = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-ALLOWED_EXTENSIONS = {"pdf", "docx", "pptx", "xlsx"}
+tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
+reader = easyocr.Reader(["en"])  
+
+ALLOWED_EXTENSIONS = {"pdf", "docx", "pptx", "xlsx", "png", "jpg", "jpeg"}
 
-# ✅ Function to Validate File Type
 def validate_file_type(file):
     if file is None:
         return "❌ No file uploaded!"
-    if isinstance(file, str):  # Text-based input (NamedString)
+    if isinstance(file, str):  
         return None
     if hasattr(file, "name"):
         ext = file.name.split(".")[-1].lower()
@@ -160,23 +162,20 @@ def validate_file_type(file):
         return None
     return "❌ Invalid file format!"
 
-# ✅ Extract Text from PDF
 # ✅ Extract Text from PDF
 def extract_text_from_pdf(file_bytes):
     try:
-        doc = pymupdf.open(stream=file_bytes, filetype="pdf")  # Use pymupdf.open()
-        return "\n".join([page.get_text() for page in doc])
+        doc = pymupdf.open(stream=file_bytes, filetype="pdf")
+        return "\n".join([page.get_text("text") for page in doc])
     except Exception as e:
-        print(f"❌ PDF Extraction Error: {e}")  # Log error
         return f"❌ PDF Error: {str(e)}"
 
 # ✅ Extract Text from DOCX & PPTX using Tika
 def extract_text_with_tika(file_bytes):
     try:
         parsed = parser.from_buffer(file_bytes)
-        return parsed["content"]
+        return parsed.get("content", "⚠️ No text found.").strip()
     except Exception as e:
-        print(f"❌ Tika Extraction Error: {e}")  # Log error
         return f"❌ Tika Error: {str(e)}"
 
 # ✅ Extract Text from Excel
@@ -187,14 +186,41 @@ def extract_text_from_excel(file_bytes):
         for sheet in wb.worksheets:
             for row in sheet.iter_rows(values_only=True):
                 text.append(" ".join(str(cell) for cell in row if cell))
-        return "\n".join(text)
+        return "\n".join(text) if text else "⚠️ No text found."
     except Exception as e:
-        print(f"❌ Excel Extraction Error: {e}")  # Log error
         return f"❌ Excel Error: {str(e)}"
 
+# ✅ Extract Text from PPTX
+def extract_text_from_pptx(file_bytes):
+    try:
+        ppt = Presentation(BytesIO(file_bytes))
+        text = []
+        for slide in ppt.slides:
+            for shape in slide.shapes:
+                if hasattr(shape, "text"):
+                    text.append(shape.text)
+        return "\n".join(text) if text else "⚠️ No text found."
+    except Exception as e:
+        return f"❌ PPTX Error: {str(e)}"
+
+# ✅ Extract Text from Image using OCR
+def extract_text_from_image(image_file):
+    try:
+        image = Image.open(image_file).convert("RGB")
+        np_image = np.array(image)
+        
+        if np_image.std() < 10:  # Low contrast check
+            return "⚠️ No meaningful content detected in the image."
+        
+        result = reader.readtext(np_image)
+        return " ".join([res[1] for res in result]) if result else "⚠️ No text found."
+    except Exception as e:
+        return f"❌ Image OCR Error: {str(e)}"
+
 # ✅ Truncate Long Text for Model
-def truncate_text(text, max_length=2048):
-    return text[:max_length] if len(text) > max_length else text
+def truncate_text(text, max_tokens=450):
+    tokens = tokenizer.tokenize(text)
+    return tokenizer.convert_tokens_to_string(tokens[:max_tokens])
 
 # ✅ Answer Questions from Image or Document
 def answer_question(file, question: str):
@@ -211,24 +237,19 @@ def answer_question(file, question: str):
         if validation_error:
             return validation_error
 
-        # ✅ Determine File Path or Read Bytes
+        # ✅ Read File Bytes
         file_bytes = None
         file_ext = None
 
-        if isinstance(file, str):  # Gradio sometimes passes a file path string
-            if os.path.exists(file):  # If it's a valid file path
-                file_ext = file.split(".")[-1].lower()
-                with open(file, "rb") as f:
-                    file_bytes = f.read()
-            else:
-                return f"❌ Error: File path does not exist! Path: {file}"
-
-        elif hasattr(file, "read"):  # If it's a file-like object
+        if isinstance(file, str) and os.path.exists(file):
+            file_ext = file.split(".")[-1].lower()
+            with open(file, "rb") as f:
+                file_bytes = f.read()
+        elif hasattr(file, "read"):
             file_ext = file.name.split(".")[-1].lower() if hasattr(file, "name") else None
             file_bytes = file.read()
-
         else:
-            return f"❌ Unexpected file type received! Type: {type(file)}"
+            return "❌ Unexpected file type received!"
 
         # ✅ Extract Text Based on File Type
         if file_ext == "pdf":
@@ -237,11 +258,12 @@ def answer_question(file, question: str):
             text = extract_text_with_tika(file_bytes)
         elif file_ext == "xlsx":
             text = extract_text_from_excel(file_bytes)
+        elif file_ext in ["png", "jpg", "jpeg"]:
+            text = extract_text_from_image(BytesIO(file_bytes))
         else:
             return f"❌ Unsupported file format: {file_ext}"
 
-        # ✅ Validate Extraction
-        if not text or "❌" in text:
+        if not text or "⚠️" in text:
             return f"⚠️ No text extracted. Error: {text}"
 
         truncated_text = truncate_text(text)
@@ -250,19 +272,15 @@ def answer_question(file, question: str):
         return response[0]["generated_text"]
 
     except Exception as e:
-        print(f"❌ General Processing Error: {e}")  # Log error to console
         return f"❌ Processing Error: {str(e)}"
 
-# ✅ Gradio Interface (Unified for Images & Documents)
+# ✅ Gradio Interface
 with gr.Blocks() as demo:
     gr.Markdown("## 📄 AI-Powered Document & Image QA")
-
     with gr.Row():
         file_input = gr.File(label="Upload Document / Image")
         question_input = gr.Textbox(label="Ask a Question", placeholder="What is this document about?")
-    
     answer_output = gr.Textbox(label="Answer")
-    
     submit_btn = gr.Button("Get Answer")
     submit_btn.click(answer_question, inputs=[file_input, question_input], outputs=answer_output)
 
@@ -273,6 +291,5 @@ app = gr.mount_gradio_app(app, demo, path="/")
 def home():
     return RedirectResponse(url="/")
 
-# ✅ Run FastAPI + Gradio
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)