Spaces:
Running
Running
| from fastapi import FastAPI, File, UploadFile | |
| import pdfplumber | |
| import docx | |
| 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 | |
| # Initialize FastAPI | |
| app = FastAPI() | |
| # Load AI Model for Question Answering | |
| qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-large", tokenizer="google/flan-t5-large", use_fast=True) | |
| # Load Pretrained Object Detection Model (Torchvision) | |
| model = fasterrcnn_resnet50_fpn(pretrained=True) | |
| model.eval() | |
| # Image Transformations | |
| transform = transforms.Compose([ | |
| transforms.ToTensor() | |
| ]) | |
| # 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 | |
| def extract_text_from_pdf(pdf_file): | |
| text = "" | |
| with pdfplumber.open(pdf_file) as pdf: | |
| for page in pdf.pages: | |
| text += page.extract_text() + "\n" | |
| return text.strip() | |
| def extract_text_from_docx(docx_file): | |
| doc = docx.Document(docx_file) | |
| return "\n".join([para.text for para in doc.paragraphs]) | |
| def extract_text_from_pptx(pptx_file): | |
| 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) | |
| def extract_text_from_excel(excel_file): | |
| wb = openpyxl.load_workbook(excel_file) | |
| text = [] | |
| for sheet in wb.worksheets: | |
| for row in sheet.iter_rows(values_only=True): | |
| text.append(" ".join(map(str, row))) | |
| return "\n".join(text) | |
| # Function to perform object detection using Torchvision | |
| def extract_text_from_image(image_file): | |
| if isinstance(image_file, np.ndarray): # Check if input is a NumPy array | |
| image = Image.fromarray(image_file) # Convert NumPy array to PIL image | |
| else: | |
| image = Image.open(image_file).convert("RGB") # Handle file input | |
| reader = easyocr.Reader(["en"]) | |
| result = reader.readtext(np.array(image)) # Convert PIL image back to NumPy array | |
| 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() | |
| if file_ext == "pdf": | |
| text = extract_text_from_pdf(file) | |
| elif file_ext == "docx": | |
| text = extract_text_from_docx(file) | |
| elif file_ext == "pptx": | |
| text = extract_text_from_pptx(file) | |
| elif file_ext == "xlsx": | |
| text = extract_text_from_excel(file) | |
| else: | |
| return "Unsupported file format!" | |
| 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"] | |
| # 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 meaningful content 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 & Image QA | |
| doc_interface = gr.Interface( | |
| fn=answer_question_from_document, | |
| inputs=[gr.File(label="Upload Document"), gr.Textbox(label="Ask a Question")], | |
| outputs="text", | |
| title="AI Document Question Answering" | |
| ) | |
| 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="/") | |
| def home(): | |
| return RedirectResponse(url="/") | |