Spaces:
Running
Running
File size: 3,340 Bytes
8bf603d b441fd7 8bf603d b441fd7 8bf603d b441fd7 8bf603d b441fd7 8bf603d b441fd7 8bf603d b441fd7 8bf603d b441fd7 8bf603d b441fd7 8bf603d b441fd7 8bf603d b441fd7 8bf603d b441fd7 8bf603d b441fd7 8bf603d b441fd7 8bf603d b441fd7 8bf603d b441fd7 8bf603d b441fd7 8bf603d b441fd7 8bf603d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | 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() |