Update app.py

app.py

@@ -4,29 +4,47 @@ from tika import parser  # Apache Tika for document parsing
 import openpyxl
 from pptx import Presentation
 import torch
+from torchvision import transforms
+from torchvision.models.detection import fasterrcnn_resnet50_fpn
 from PIL import Image
 from transformers import pipeline
 import gradio as gr
+from fastapi.responses import RedirectResponse
 import numpy as np
 import easyocr
 
-# Initialize FastAPI (not needed for HF Spaces, but kept for flexibility)
+# Initialize FastAPI
 app = FastAPI()
 
+# Load AI Model for Question Answering (DeepSeek-V2-Chat)
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Preload Hugging Face model
 print(f"🔄 Loading models")
 
-doc_qa_pipeline = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", device=-1)
-image_captioning_pipeline = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
-print("✅ Models loaded")
+qa_pipeline =  pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", device=-1)
+# Load Pretrained Object Detection Model (Torchvision)
+from torchvision.models.detection import FasterRCNN_ResNet50_FPN_Weights
+weights = FasterRCNN_ResNet50_FPN_Weights.DEFAULT
+model = fasterrcnn_resnet50_fpn(weights=weights)
+model.eval()
+# Load Pretrained Object Detection Model (if needed)
+model = fasterrcnn_resnet50_fpn(pretrained=True)
+model.eval()
+
+# Initialize OCR Model (Lazy Load)
+reader = easyocr.Reader(["en"], gpu=True)
 
-# Initialize OCR Model (CPU Mode)
-reader = easyocr.Reader(["en"], gpu=False)
+# Image Transformations
+transform = transforms.Compose([
+    transforms.ToTensor()
+])
 
 # Allowed File Extensions
 ALLOWED_EXTENSIONS = {"pdf", "docx", "pptx", "xlsx"}
 
 def validate_file_type(file):
-    ext = file.filename.split(".")[-1].lower()
+    ext = file.name.split(".")[-1].lower()
     print(f"🔍 Validating file type: {ext}")
     if ext not in ALLOWED_EXTENSIONS:
         return f"❌ Unsupported file format: {ext}"
@@ -40,27 +58,40 @@ def truncate_text(text, max_tokens=450):
     return truncated
 
 # Document Text Extraction Functions
-def extract_text_from_pdf(pdf_bytes):
+def extract_text_from_pdf(pdf_file):
     try:
         print("📄 Extracting text from PDF...")
-        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
+        doc = fitz.open(pdf_file)
         text = "\n".join([page.get_text("text") for page in doc])
         return text if text else "⚠️ No text found."
     except Exception as e:
         return f"❌ Error reading PDF: {str(e)}"
 
-def extract_text_with_tika(file_bytes):
+def extract_text_with_tika(file):
     try:
         print("📝 Extracting text with Tika...")
-        parsed = parser.from_buffer(file_bytes)
+        parsed = parser.from_buffer(file)
         return parsed.get("content", "⚠️ No text found.").strip()
     except Exception as e:
         return f"❌ Error reading document: {str(e)}"
 
-def extract_text_from_excel(excel_bytes):
+def extract_text_from_pptx(pptx_file):
+    try:
+        print("📊 Extracting text from PPTX...")
+        ppt = Presentation(pptx_file)
+        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"❌ Error reading PPTX: {str(e)}"
+
+def extract_text_from_excel(excel_file):
     try:
         print("📊 Extracting text from Excel...")
-        wb = openpyxl.load_workbook(excel_bytes, read_only=True)
+        wb = openpyxl.load_workbook(excel_file, read_only=True)
         text = []
         for sheet in wb.worksheets:
             for row in sheet.iter_rows(values_only=True):
@@ -69,21 +100,29 @@ def extract_text_from_excel(excel_bytes):
     except Exception as e:
         return f"❌ Error reading Excel: {str(e)}"
 
-def answer_question_from_document(file: UploadFile, question: str):
+def extract_text_from_image(image_file):
+    print("🖼️ Extracting text from image...")
+    image = Image.open(image_file).convert("RGB")
+    if np.array(image).std() < 10:  # Low contrast = likely empty
+        return "⚠️ No meaningful content detected in the image."
+    
+    result = reader.readtext(np.array(image))
+    return " ".join([res[1] for res in result]) if result else "⚠️ No text found."
+
+# Function to answer questions based on document content
+def answer_question_from_document(file, question):
     print("📂 Processing document for QA...")
     validation_error = validate_file_type(file)
     if validation_error:
         return validation_error
     
-    file_ext = file.filename.split(".")[-1].lower()
-    file_bytes = file.file.read()
-
+    file_ext = file.name.split(".")[-1].lower()
     if file_ext == "pdf":
-        text = extract_text_from_pdf(file_bytes)
+        text = extract_text_from_pdf(file)
     elif file_ext in ["docx", "pptx"]:
-        text = extract_text_with_tika(file_bytes)
+        text = extract_text_with_tika(file)
     elif file_ext == "xlsx":
-        text = extract_text_from_excel(file_bytes)
+        text = extract_text_from_excel(file)
     else:
         return "❌ Unsupported file format!"
     
@@ -92,25 +131,21 @@ def answer_question_from_document(file: UploadFile, question: str):
     
     truncated_text = truncate_text(text)
     print("🤖 Generating response...")
-    response = doc_qa_pipeline(f"Question: {question}\nContext: {truncated_text}")
+    response = qa_pipeline(f"Question: {question}\nContext: {truncated_text}")
     
     return response[0]["generated_text"]
 
 def answer_question_from_image(image, question):
-    try:
-        print("🖼️ Processing image for QA...")
-        if isinstance(image, np.ndarray):  # If it's a NumPy array from Gradio
-            image = Image.fromarray(image)  # Convert to PIL Image
-        
-        print("🖼️ Generating caption for image...")
-        caption = image_captioning_pipeline(image)[0]['generated_text']
-
-        print("🤖 Answering question based on caption...")
-        response = doc_qa_pipeline(f"Question: {question}\nContext: {caption}")
-
-        return response[0]["generated_text"]
-    except Exception as e:
-        return f"❌ Error processing image: {str(e)}"
+    print("🖼️ Processing image for QA...")
+    image_text = extract_text_from_image(image)
+    if not image_text:
+        return "⚠️ No meaningful content detected in the image."
+    
+    truncated_text = truncate_text(image_text)
+    print("🤖 Generating response...")
+    response = qa_pipeline(f"Question: {question}\nContext: {truncated_text}")
+    
+    return response[0]["generated_text"]
 
 # Gradio UI for Document & Image QA
 doc_interface = gr.Interface(
@@ -127,8 +162,10 @@ img_interface = gr.Interface(
     title="🖼️ AI Image Question Answering"
 )
 
-# Launch Gradio
-app = gr.TabbedInterface([doc_interface, img_interface], ["📄 Document QA", "🖼️ Image QA"])
+# Mount Gradio Interfaces
+demo = gr.TabbedInterface([doc_interface, img_interface], ["📄 Document QA", "🖼️ Image QA"])
+app = gr.mount_gradio_app(app, demo, path="/")
 
-if __name__ == "__main__":
-    app.launch(share=True)  # For Hugging Face Spaces
+@app.get("/")
+def home():
+    return RedirectResponse(url="/")