Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| from transformers import NougatProcessor, VisionEncoderDecoderModel | |
| from PIL import Image | |
| import tempfile | |
| import os | |
| import re | |
| import fitz # PyMuPDF | |
| from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer | |
| from reportlab.lib.styles import getSampleStyleSheet | |
| from reportlab.lib.units import inch | |
| # ----------------------------- | |
| # Load Model | |
| # ----------------------------- | |
| MODEL_NAME = "facebook/nougat-base" | |
| processor = NougatProcessor.from_pretrained(MODEL_NAME) | |
| model = VisionEncoderDecoderModel.from_pretrained(MODEL_NAME) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model.to(device) | |
| # ----------------------------- | |
| # Utilities | |
| # ----------------------------- | |
| def clean_output(text): | |
| text = re.sub(r"<.*?>", "", text) | |
| return text.strip() | |
| def image_to_latex(image): | |
| pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(device) | |
| outputs = model.generate(pixel_values, max_new_tokens=1024) | |
| latex = processor.batch_decode(outputs, skip_special_tokens=True)[0] | |
| return clean_output(latex) | |
| def pdf_to_images(pdf_path): | |
| images = [] | |
| pdf = fitz.open(pdf_path) | |
| for page in pdf: | |
| pix = page.get_pixmap() | |
| img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) | |
| images.append(img) | |
| return images | |
| def generate_pdf_from_text(text, output_path): | |
| doc = SimpleDocTemplate(output_path) | |
| elements = [] | |
| styles = getSampleStyleSheet() | |
| for line in text.split("\n"): | |
| elements.append(Paragraph(line, styles["Normal"])) | |
| elements.append(Spacer(1, 0.2 * inch)) | |
| doc.build(elements) | |
| # ----------------------------- | |
| # Main Function | |
| # ----------------------------- | |
| def process_file(file): | |
| if file is None: | |
| return "Upload file first", None, None | |
| try: | |
| temp_dir = tempfile.mkdtemp() | |
| full_latex = "" | |
| if file.name.endswith(".pdf"): | |
| images = pdf_to_images(file.name) | |
| for img in images: | |
| latex = image_to_latex(img) | |
| full_latex += latex + "\n\n" | |
| else: | |
| image = Image.open(file.name).convert("RGB") | |
| full_latex = image_to_latex(image) | |
| # Save .tex file | |
| tex_path = os.path.join(temp_dir, "output.tex") | |
| with open(tex_path, "w", encoding="utf-8") as f: | |
| f.write(full_latex) | |
| # Generate readable PDF | |
| pdf_path = os.path.join(temp_dir, "output.pdf") | |
| generate_pdf_from_text(full_latex, pdf_path) | |
| return full_latex, tex_path, pdf_path | |
| except Exception as e: | |
| return f"Error occurred: {str(e)}", None, None | |
| # ----------------------------- | |
| # Gradio UI | |
| # ----------------------------- | |
| with gr.Blocks(title="Handwritten Assignment Digitizer") as demo: | |
| gr.Markdown(""" | |
| # Handwritten Assignment Digitizer | |
| """) | |
| file_input = gr.File(label="Upload Image or PDF") | |
| latex_output = gr.Textbox(label="Generated LaTeX Code", lines=15) | |
| tex_download = gr.File(label="Download .tex file") | |
| pdf_download = gr.File(label="Download Readable PDF") | |
| process_button = gr.Button("Convert Document") | |
| process_button.click( | |
| fn=process_file, | |
| inputs=file_input, | |
| outputs=[latex_output, tex_download, pdf_download] | |
| ) | |
| demo.launch() |