Update app.py

app.py

@@ -1,253 +1,33 @@
-"""
-ScriptFlow — Handwriting PDF → Word Document Converter
-Hugging Face Spaces (CPU Free Tier)
-"""
-
-import os
-import glob
-import shutil
-import tempfile
-import subprocess
-from pathlib import Path
-
 import gradio as gr
-import torch
 from PIL import Image
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel, pipeline
-from docx import Document
-from docx.shared import Pt, Inches
-from docx.enum.text import WD_ALIGN_PARAGRAPH
-
-
-# ──────────────────────────────────────────────────────────────────────────────
-# Model Caching Setup
-# ──────────────────────────────────────────────────────────────────────────────
-
-_trocr_cache = {}
-_grammar_pipe = None
-
-
-def get_trocr_model(model_name: str):
-    if model_name not in _trocr_cache:
-        processor = TrOCRProcessor.from_pretrained(model_name)
-        model = VisionEncoderDecoderModel.from_pretrained(model_name)
-        model.eval()
-        _trocr_cache[model_name] = (processor, model)
-    return _trocr_cache[model_name]
-
-
-def get_grammar_pipeline():
-    global _grammar_pipe
-    if _grammar_pipe is None:
-        _grammar_pipe = pipeline(
-            "text2text-generation",
-            model="vennify/t5-base-grammar-correction",
-            tokenizer="vennify/t5-base-grammar-correction",
-        )
-    return _grammar_pipe
-
-
-# ──────────────────────────────────────────────────────────────────────────────
-# Core Processing Functions
-# ──────────────────────────────────────────────────────────────────────────────
-
-def pdf_to_images(pdf_path: str, out_dir: str) -> list[str]:
-    pages_dir = os.path.join(out_dir, "pages")
-    os.makedirs(pages_dir, exist_ok=True)
-
-    result = subprocess.run(
-        ["pdftoppm", "-png", "-r", "200", pdf_path, os.path.join(pages_dir, "page")],
-        capture_output=True,
-        text=True,
-        timeout=180,
-    )
-    if result.returncode != 0:
-        raise RuntimeError(f"pdftoppm failed: {result.stderr.strip()}")
-
-    images = sorted(glob.glob(os.path.join(pages_dir, "*.png")))
-    if not images:
-        raise RuntimeError("No page images could be extracted from the PDF.")
-    return images
-
-
-def ocr_images(image_paths: list[str], model_name: str, progress: gr.Progress) -> list[str]:
-    processor, model = get_trocr_model(model_name)
-    page_texts = []
-
-    total_pages = len(image_paths)
-    for idx, img_path in enumerate(image_paths):
-        progress((0.1 + (idx / total_pages) * 0.5), desc=f"OCR processing: page {idx + 1}/{total_pages}")
-        
-        image = Image.open(img_path).convert("RGB")
-        pixel_values = processor(images=image, return_tensors="pt").pixel_values
-        
-        with torch.no_grad():
-            generated_ids = model.generate(pixel_values)
-        
-        text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        page_texts.append(text.strip())
-
-    return page_texts
-
-
-def refine_grammar(raw_pages: list[str], progress: gr.Progress) -> list[str]:
-    pipe = get_grammar_pipeline()
-    refined_pages = []
-
-    total_pages = len(raw_pages)
-    for idx, page in enumerate(raw_pages):
-        progress((0.6 + (idx / total_pages) * 0.3), desc=f"Refining grammar: page {idx + 1}/{total_pages}")
-        
-        if not page.strip():
-            refined_pages.append("")
-            continue
-
-        # Chunk text into 400-char segments for T5 context window
-        chunks = [page[i:i + 400] for i in range(0, len(page), 400)]
-        corrected_chunks = []
-        
-        for chunk in chunks:
-            if not chunk.strip():
-                continue
-            out = pipe(
-                f"grammar: {chunk}",
-                max_length=512,
-                num_beams=4,
-                early_stopping=True,
-            )
-            corrected_chunks.append(out[0]["generated_text"])
-
-        refined_pages.append(" ".join(corrected_chunks) if corrected_chunks else page)
-
-    return refined_pages
-
-
-def export_to_docx(pages: list[str], output_path: str):
-    doc = Document()
-
-    for section in doc.sections:
-        section.top_margin = Inches(1)
-        section.bottom_margin = Inches(1)
-        section.left_margin = Inches(1)
-        section.right_margin = Inches(1)
-
-    heading = doc.add_heading("Converted Handwritten Document", level=1)
-    heading.alignment = WD_ALIGN_PARAGRAPH.CENTER
-    doc.add_paragraph()
-
-    for idx, page_text in enumerate(pages):
-        if idx > 0:
-            doc.add_page_break()
-
-        lines = page_text.split("\n")
-        for line in lines:
-            line_str = line.strip()
-            p = doc.add_paragraph()
-            if line_str:
-                run = p.add_run(line_str)
-                run.font.name = "Calibri"
-                run.font.size = Pt(12)
-
-    doc.save(output_path)
-
-
-# ──────────────────────────────────────────────────────────────────────────────
-# Main Pipeline Executable
-# ──────────────────────────────────────────────────────────────────────────────
-
-def convert_pdf(pdf_file, trocr_model: str, use_grammar: bool, progress=gr.Progress()):
-    if pdf_file is None:
-        raise gr.Error("Please upload a handwritten PDF file.")
-
-    # Cross-version filepath handling
-    pdf_path = pdf_file if isinstance(pdf_file, str) else getattr(pdf_file, "name", str(pdf_file))
-
-    with tempfile.TemporaryDirectory() as tmp_dir:
-        # Step 1: Convert PDF to PNGs
-        progress(0.05, desc="Extracting pages from PDF...")
-        images = pdf_to_images(pdf_path, tmp_dir)
-
-        # Step 2: OCR with TrOCR
-        progress(0.10, desc="Loading OCR Model...")
-        raw_pages = ocr_images(images, trocr_model, progress)
-        raw_text_combined = "\n\n--- [PAGE BREAK] ---\n\n".join(raw_pages)
-
-        # Step 3: Optional T5 Grammar Correction
-        if use_grammar:
-            progress(0.60, desc="Loading Grammar Model...")
-            final_pages = refine_grammar(raw_pages, progress)
-        else:
-            final_pages = raw_pages
-
-        final_text_combined = "\n\n--- [PAGE BREAK] ---\n\n".join(final_pages)
-
-        # Step 4: Write .docx File
-        progress(0.92, desc="Creating Word Document...")
-        docx_temp_path = os.path.join(tmp_dir, "ScriptFlow_Output.docx")
-        export_to_docx(final_pages, docx_temp_path)
-
-        # Export file out of temporary scope for downloading
-        output_file = tempfile.NamedTemporaryFile(suffix=".docx", delete=False, prefix="scriptflow_")
-        output_file.close()
-        shutil.copy(docx_temp_path, output_file.name)
-
-    progress(1.0, desc="Conversion Complete!")
-    return raw_text_combined, final_text_combined, output_file.name
-
-
-# ──────────────────────────────────────────────────────────────────────────────
-# Gradio UI Design
-# ──────────────────────────────────────────────────────────────────────────────
-
-with gr.Blocks(title="ScriptFlow — PDF to Word Converter", theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # ✍️ ScriptFlow
-        ### Convert Handwritten PDFs to Word (.docx) Documents
-        *Extract handwritten text with TrOCR, clean up syntax with T5, and download a ready-to-edit `.docx` file.*
-        """
-    )
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            pdf_input = gr.File(
-                label="Upload Handwritten PDF",
-                file_types=[".pdf"],
-                type="filepath",
-            )
-            model_selector = gr.Radio(
-                choices=[
-                    ("Small (~330 MB) — Faster", "microsoft/trocr-small-handwritten"),
-                    ("Base (~440 MB) — More Accurate", "microsoft/trocr-base-handwritten"),
-                ],
-                value="microsoft/trocr-small-handwritten",
-                label="TrOCR Model",
-            )
-            grammar_toggle = gr.Checkbox(
-                label="Enable T5 Grammar & Refinement",
-                value=False,
-            )
-            convert_btn = gr.Button("Convert Document", variant="primary")
-
-        with gr.Column(scale=2):
-            raw_output = gr.Textbox(label="Raw OCR Output", lines=10, interactive=False)
-            refined_output = gr.Textbox(label="Refined Output", lines=10, interactive=False)
-            docx_download = gr.File(label="Download Word (.docx) File", interactive=False)
-
-    gr.Markdown(
-        """
-        ---
-        **Tips for best results:**
-        * High-resolution scans (200+ DPI) work significantly better.
-        * Grammar correction adds extra processing time per page on CPU.
-        """
-    )
-
-    convert_btn.click(
-        fn=convert_pdf,
-        inputs=[pdf_input, model_selector, grammar_toggle],
-        outputs=[raw_output, refined_output, docx_download],
-    )
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+
+# Load lightweight CPU-friendly model
+processor = TrOCRProcessor.from_pretrained("microsoft/trocr-small-handwritten")
+model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-small-handwritten")
+
+def parse_handwriting(image):
+    if image is None:
+        return "Please upload an image."
+    
+    # Convert image to RGB
+    image = image.convert("RGB")
+    
+    # Process image and generate text
+    pixel_values = processor(images=image, return_tensors="pt").pixel_values
+    generated_ids = model.generate(pixel_values)
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    
+    return generated_text
+
+# Simple Gradio interface
+demo = gr.Interface(
+    fn=parse_handwriting,
+    inputs=gr.Image(type="pil", label="Upload Handwritten Image"),
+    outputs=gr.Textbox(label="Extracted Text"),
+    title="Handwriting Parser (CPU Friendly)",
+    description="Upload a cropped image of handwritten text to transcribe it."
+)
 
 if __name__ == "__main__":
     demo.launch()