Spaces:
Running
Running
| from fastapi import FastAPI, UploadFile, File | |
| from fastapi.responses import RedirectResponse | |
| import gradio as gr | |
| from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM | |
| import tempfile | |
| import os | |
| from PIL import Image | |
| import fitz # PyMuPDF | |
| import docx | |
| import openpyxl | |
| from pptx import Presentation | |
| import easyocr | |
| app = FastAPI() | |
| # Initialize models with error handling | |
| try: | |
| # Load summarization model directly with tokenizer | |
| tokenizer = AutoTokenizer.from_pretrained("FeruzaBoynazarovaas/my_awesome_billsum_model", use_fast=False) | |
| model = AutoModelForSeq2SeqLM.from_pretrained("FeruzaBoynazarovaas/my_awesome_billsum_model") | |
| summarizer = pipeline( | |
| "text2text-generation", | |
| model=model, | |
| tokenizer=tokenizer | |
| ) | |
| except Exception as e: | |
| print(f"Error loading summarizer: {e}") | |
| # Fallback to a default model if custom fails | |
| summarizer = pipeline("text2text-generation", model="t5-small") | |
| # Other models (these should work fine) | |
| captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning") | |
| reader = easyocr.Reader(['en']) | |
| def extract_text_from_file(file_path: str, file_type: str): | |
| """Extract text from different document formats""" | |
| try: | |
| if file_type == "pdf": | |
| with fitz.open(file_path) as doc: | |
| return "\n".join(page.get_text() for page in doc) | |
| elif file_type == "docx": | |
| doc = docx.Document(file_path) | |
| return "\n".join(p.text for p in doc.paragraphs) | |
| elif file_type == "pptx": | |
| prs = Presentation(file_path) | |
| return "\n".join(shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, "text")) | |
| elif file_type == "xlsx": | |
| wb = openpyxl.load_workbook(file_path) | |
| return "\n".join(str(cell.value) for sheet in wb for row in sheet for cell in row) | |
| else: | |
| return "Unsupported file format" | |
| except Exception as e: | |
| return f"Error reading file: {str(e)}" | |
| def process_document(file): | |
| try: | |
| file_ext = os.path.splitext(file.name)[1][1:].lower() | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_ext}") as tmp: | |
| tmp.write(file.read()) | |
| tmp_path = tmp.name | |
| text = extract_text_from_file(tmp_path, file_ext) | |
| summary = summarizer(text, max_length=150, min_length=30, do_sample=False)[0]['generated_text'] | |
| os.unlink(tmp_path) | |
| return summary | |
| except Exception as e: | |
| return f"Processing error: {str(e)}" | |
| def process_image(image): | |
| try: | |
| img = Image.open(image) | |
| caption = captioner(img)[0]['generated_text'] | |
| ocr_result = reader.readtext(img) | |
| ocr_text = " ".join([res[1] for res in ocr_result]) | |
| return { | |
| "caption": caption, | |
| "ocr_text": ocr_text if ocr_text else "No readable text found" | |
| } | |
| except Exception as e: | |
| return {"error": str(e)} | |
| # Gradio Interface | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# 📄 Document & Image Analysis") | |
| with gr.Tab("Document Summarization"): | |
| doc_input = gr.File(label="Upload Document") | |
| doc_output = gr.Textbox(label="Summary") | |
| doc_button = gr.Button("Summarize") | |
| with gr.Tab("Image Analysis"): | |
| img_input = gr.Image(type="filepath", label="Upload Image") | |
| caption_output = gr.Textbox(label="Image Caption") | |
| ocr_output = gr.Textbox(label="Extracted Text") | |
| img_button = gr.Button("Analyze") | |
| doc_button.click(process_document, inputs=doc_input, outputs=doc_output) | |
| img_button.click(process_image, inputs=img_input, outputs=[caption_output, ocr_output]) | |
| app = gr.mount_gradio_app(app, demo, path="/") | |
| def redirect(): | |
| return RedirectResponse(url="/") |