Update app.py

app.py

@@ -1,141 +1,113 @@
+from fastapi import FastAPI, Form, File, UploadFile
+from fastapi.responses import RedirectResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
+from transformers import pipeline
 import os
-from fastapi import FastAPI, UploadFile, File
-import gradio as gr
-from transformers import pipeline, AutoModelForQuestionAnswering, AutoTokenizer
+from PIL import Image
+import io
 import pdfplumber
 import docx
 import openpyxl
 import pytesseract
-from PIL import Image
-from easyocr import Reader
+from io import BytesIO
+import fitz  # PyMuPDF
+import easyocr
+from fastapi.templating import Jinja2Templates
+from starlette.requests import Request
 
-# Initialize FastAPI app
+# Initialize the app
 app = FastAPI()
 
-# Load the question-answering models
-doc_qa_model_name = "microsoft/phi-2"
-image_qa_model_name = "Salesforce/blip-vqa-base"
+# Mount the static directory to serve HTML, CSS, JS files
+app.mount("/static", StaticFiles(directory="static"), name="static")
 
-# Load the models and tokenizers for each use case
-doc_qa_model = AutoModelForQuestionAnswering.from_pretrained(doc_qa_model_name)
-doc_qa_tokenizer = AutoTokenizer.from_pretrained(doc_qa_model_name)
+# Initialize transformers pipelines
+qa_pipeline = pipeline("question-answering", model="microsoft/phi-2", tokenizer="microsoft/phi-2")
+image_qa_pipeline = pipeline("image-question-answering", model="Salesforce/blip-vqa-base", tokenizer="Salesforce/blip-vqa-base")
 
-image_qa_model = AutoModelForQuestionAnswering.from_pretrained(image_qa_model_name)
-image_qa_tokenizer = AutoTokenizer.from_pretrained(image_qa_model_name)
+# Initialize EasyOCR for image-based text extraction
+reader = easyocr.Reader(['en'])
 
-# Load VQA pipeline for images
-image_qa_pipeline = pipeline("image-captioning", model=image_qa_model, tokenizer=image_qa_tokenizer)
+# Define a template for rendering HTML
+templates = Jinja2Templates(directory="templates")
 
-# Load OCR Reader
-ocr_reader = Reader(lang_list=['en'])
-
-# Helper function to process DOCX files
-def extract_text_from_docx(file_path):
-    doc = docx.Document(file_path)
-    text = "\n".join([para.text for para in doc.paragraphs])
-    return text
-
-# Helper function to process Excel files
-def extract_text_from_excel(file_path):
-    wb = openpyxl.load_workbook(file_path)
-    sheet = wb.active
-    text = ""
-    for row in sheet.iter_rows():
-        text += " ".join([str(cell.value) for cell in row]) + "\n"
-    return text
-
-# Helper function to process PDF files
-def extract_text_from_pdf(file_path):
+# Function to process PDFs
+def extract_pdf_text(file_path: str):
     with pdfplumber.open(file_path) as pdf:
         text = ""
         for page in pdf.pages:
             text += page.extract_text()
     return text
 
-# Helper function to process images (OCR)
-def extract_text_from_image(image_path):
-    image = Image.open(image_path)
-    text = pytesseract.image_to_string(image)
+# Function to process DOCX files
+def extract_docx_text(file_path: str):
+    doc = docx.Document(file_path)
+    text = ""
+    for para in doc.paragraphs:
+        text += para.text
     return text
 
-# AI-powered Question Answering for documents
-def document_question_answering(document_text, question):
-    inputs = doc_qa_tokenizer.encode(question, document_text, return_tensors='pt')
-    outputs = doc_qa_model(inputs)
-    answer_start = outputs.start_logits.argmax()
-    answer_end = outputs.end_logits.argmax()
-    answer = doc_qa_tokenizer.decode(inputs[0][answer_start:answer_end + 1])
-    return answer
+# Function to process PPTX files
+def extract_pptx_text(file_path: str):
+    from pptx import Presentation
+    prs = Presentation(file_path)
+    text = ""
+    for slide in prs.slides:
+        for shape in slide.shapes:
+            if hasattr(shape, "text"):
+                text += shape.text
+    return text
 
-# AI-powered Question Answering for images
-def image_question_answering(image, question):
-    # Generate caption for the image first
-    caption = image_qa_pipeline(image)[0]['caption']
-    # Combine caption with the question to form input to the QA model
-    return caption  # For simplicity, you may further process this if needed
+# Function to extract text from images using OCR
+def extract_text_from_image(image: Image):
+    text = pytesseract.image_to_string(image)
+    return text
 
-# FastAPI route for document-based question answering
-@app.post("/qa/document")
-async def qa_document(file: UploadFile = File(...), question: str = None):
-    file_location = f"temp_files/{file.filename}"
-    with open(file_location, "wb") as f:
+# Home route
+@app.get("/")
+def home():
+    return RedirectResponse(url="/docs")
+
+# Function to answer questions based on document content
+@app.post("/question-answering-doc")
+async def question_answering_doc(question: str = Form(...), file: UploadFile = File(...)):
+    # Save the uploaded file temporarily
+    file_path = f"temp_files/{file.filename}"
+    os.makedirs(os.path.dirname(file_path), exist_ok=True)
+    with open(file_path, "wb") as f:
         f.write(await file.read())
     
-    # Extract text from the document based on its format
-    if file.filename.endswith('.pdf'):
-        document_text = extract_text_from_pdf(file_location)
-    elif file.filename.endswith('.docx'):
-        document_text = extract_text_from_docx(file_location)
-    elif file.filename.endswith('.xlsx'):
-        document_text = extract_text_from_excel(file_location)
+    # Extract text based on file type
+    if file.filename.endswith(".pdf"):
+        text = extract_pdf_text(file_path)
+    elif file.filename.endswith(".docx"):
+        text = extract_docx_text(file_path)
+    elif file.filename.endswith(".pptx"):
+        text = extract_pptx_text(file_path)
     else:
-        return {"error": "Unsupported file format."}
+        return {"error": "Unsupported file format"}
 
-    # Get the answer using the document QA model
-    answer = document_question_answering(document_text, question)
-    return {"answer": answer}
+    # Use the model for question answering
+    qa_result = qa_pipeline(question=question, context=text)
+    return {"answer": qa_result['answer']}
 
-# FastAPI route for image-based question answering
-@app.post("/qa/image")
-async def qa_image(file: UploadFile = File(...), question: str = None):
-    file_location = f"temp_files/{file.filename}"
-    with open(file_location, "wb") as f:
-        f.write(await file.read())
+# Function to answer questions based on images
+@app.post("/question-answering-image")
+async def question_answering_image(question: str = Form(...), image_file: UploadFile = File(...)):
+    # Open the uploaded image
+    image = Image.open(BytesIO(await image_file.read()))
     
-    # Extract text from the image using OCR
-    image_text = extract_text_from_image(file_location)
-
-    # Get the answer using the image QA model (BLIP VQA)
-    image_answer = image_question_answering(file_location, question)
-    return {"answer": image_answer}
-
-# Gradio Interface for Document QA
-def document_qa_interface(file, question):
-    file_path = file.name
-    if file_path.endswith(".pdf"):
-        document_text = extract_text_from_pdf(file_path)
-    elif file_path.endswith(".docx"):
-        document_text = extract_text_from_docx(file_path)
-    elif file_path.endswith(".xlsx"):
-        document_text = extract_text_from_excel(file_path)
-    else:
-        return "Unsupported document format."
-
-    return document_question_answering(document_text, question)
-
-# Gradio Interface for Image QA
-def image_qa_interface(image, question):
-    return image_question_answering(image, question)
-
-# Gradio Web Interface
-iface_document = gr.Interface(fn=document_qa_interface, inputs=["file", "text"], outputs="text")
-iface_image = gr.Interface(fn=image_qa_interface, inputs=["image", "text"], outputs="text")
+    # Use EasyOCR to extract text if the image has textual content
+    image_text = extract_text_from_image(image)
+    
+    # Use the BLIP VQA model for question answering on the image
+    image_qa_result = image_qa_pipeline(image=image, question=question)
+    
+    return {"answer": image_qa_result['answer'], "image_text": image_text}
 
-# Run FastAPI app
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+# Serve the application in Hugging Face space
+@app.get("/docs")
+async def get_docs(request: Request):
+    return templates.TemplateResponse("index.html", {"request": request})
 
-# Launch Gradio Interface
-iface_document.launch(share=True)
-iface_image.launch(share=True)