Update app.py

app.py

@@ -1,25 +1,54 @@
 """
-ScriptFlow — Handwriting PDF → Word Document converter
-HuggingFace Spaces (CPU free tier)
-
-Pipeline:
-  1. PDF  → page images  (pdftoppm / Poppler)
-  2. Images → raw text   (TrOCR small-handwritten)
-  3. Raw text → refined  (T5 grammar correction, optional)
-  4. Refined text → .docx (python-docx)
+ScriptFlow — Handwriting PDF → Word Document Converter
+Hugging Face Spaces (CPU Free Tier)
 """
 
 import os
 import glob
-import subprocess
-import tempfile
 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
+
 
 # ──────────────────────────────────────────────────────────────────────────────
-# Step 1 — PDF → images
+# 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]:
@@ -27,85 +56,56 @@ def pdf_to_images(pdf_path: str, out_dir: str) -> list[str]:
     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=120,
+        ["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 were generated from the PDF.")
+        raise RuntimeError("No page images could be extracted from the PDF.")
     return images
 
 
-# ──────────────────────────────────────────────────────────────────────────────
-# Step 2 — Images → raw text (TrOCR)
-# ──────────────────────────────────────────────────────────────────────────────
-
-_trocr_cache = {}
-
-def load_trocr(model_name: str):
-    global _trocr_cache
-    if model_name not in _trocr_cache:
-        from transformers import TrOCRProcessor, VisionEncoderDecoderModel
-        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 images_to_text(image_paths: list[str], model_name: str,
-                    progress: gr.Progress) -> list[str]:
-    import torch
-    from PIL import Image
-
-    processor, model = load_trocr(model_name)
+def ocr_images(image_paths: list[str], model_name: str, progress: gr.Progress) -> list[str]:
+    processor, model = get_trocr_model(model_name)
     page_texts = []
 
-    for i, img_path in enumerate(image_paths):
-        progress(i / len(image_paths),
-                 desc=f"OCR — page {i + 1} / {len(image_paths)}")
+    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():
-            ids = model.generate(pixel_values)
-        text = processor.batch_decode(ids, skip_special_tokens=True)[0]
+            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
 
 
-# ──────────────────────────────────────────────────────────────────────────────
-# Step 3 — Grammar correction (T5, optional)
-# ──────────────────────────────────────────────────────────────────────────────
-
-_grammar_pipe = None
-
-def load_grammar_pipeline():
-    global _grammar_pipe
-    if _grammar_pipe is None:
-        from transformers import pipeline
-        _grammar_pipe = pipeline(
-            "text2text-generation",
-            model="vennify/t5-base-grammar-correction",
-            tokenizer="vennify/t5-base-grammar-correction",
-        )
-    return _grammar_pipe
-
+def refine_grammar(raw_pages: list[str], progress: gr.Progress) -> list[str]:
+    pipe = get_grammar_pipeline()
+    refined_pages = []
 
-def refine_text(raw_pages: list[str], progress: gr.Progress) -> list[str]:
-    pipe = load_grammar_pipeline()
-    refined = []
+    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
 
-    for i, page in enumerate(raw_pages):
-        progress(i / len(raw_pages),
-                 desc=f"Grammar correction — page {i + 1} / {len(raw_pages)}")
-        # T5 grammar model has a 512-token limit; chunk long pages
-        chunks = [page[j:j+400] for j in range(0, len(page), 400)]
+        # 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
@@ -116,147 +116,137 @@ def refine_text(raw_pages: list[str], progress: gr.Progress) -> list[str]:
                 early_stopping=True,
             )
             corrected_chunks.append(out[0]["generated_text"])
-        refined.append(" ".join(corrected_chunks) if corrected_chunks else page)
-
-    return refined
 
+        refined_pages.append(" ".join(corrected_chunks) if corrected_chunks else page)
 
-# ──────────────────────────────────────────────────────────────────────────────
-# Step 4 — Text → .docx
-# ──────────────────────────────────────────────────────────────────────────────
+    return refined_pages
 
-def text_to_docx(pages: list[str], output_path: str):
-    from docx import Document
-    from docx.shared import Pt, Inches
-    from docx.enum.text import WD_ALIGN_PARAGRAPH
 
+def export_to_docx(pages: list[str], output_path: str):
     doc = Document()
+
     for section in doc.sections:
-        section.top_margin    = Inches(1)
+        section.top_margin = Inches(1)
         section.bottom_margin = Inches(1)
-        section.left_margin   = Inches(1.25)
-        section.right_margin  = Inches(1.25)
+        section.left_margin = Inches(1)
+        section.right_margin = Inches(1)
 
-    title = doc.add_heading("Converted Handwritten Document", level=1)
-    title.alignment = WD_ALIGN_PARAGRAPH.CENTER
-    doc.add_paragraph("")
+    heading = doc.add_heading("Converted Handwritten Document", level=1)
+    heading.alignment = WD_ALIGN_PARAGRAPH.CENTER
+    doc.add_paragraph()
 
-    for i, page_text in enumerate(pages):
-        if i > 0:
+    for idx, page_text in enumerate(pages):
+        if idx > 0:
             doc.add_page_break()
-        for line in page_text.split("\n"):
-            line = line.strip()
+
+        lines = page_text.split("\n")
+        for line in lines:
+            line_str = line.strip()
             p = doc.add_paragraph()
-            if line:
-                run = p.add_run(line)
+            if line_str:
+                run = p.add_run(line_str)
+                run.font.name = "Calibri"
                 run.font.size = Pt(12)
 
     doc.save(output_path)
 
 
 # ───────────��──────────────────────────────────────────────────────────────────
-# Main pipeline function (called by Gradio)
+# Main Pipeline Executable
 # ──────────────────────────────────────────────────────────────────────────────
 
-def run_pipeline(
-    pdf_file,
-    trocr_model: str,
-    use_grammar: bool,
-    progress=gr.Progress(track_tqdm=True),
-):
+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 PDF file first.")
+        raise gr.Error("Please upload a handwritten PDF file.")
 
-    # Handle Gradio filepath string vs object compatibility
-    pdf_path = pdf_file if isinstance(pdf_file, str) else pdf_file.name
+    # 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:
-        # Step 1
-        progress(0, desc="Converting PDF to images…")
-        images = pdf_to_images(pdf_path, tmp)
+    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
-        progress(0.1, desc="Loading TrOCR model…")
-        raw_pages = images_to_text(images, trocr_model, progress)
-        raw_text = "\n\n[PAGE BREAK]\n\n".join(raw_pages)
+        # 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
+        # Step 3: Optional T5 Grammar Correction
         if use_grammar:
-            progress(0.6, desc="Loading grammar correction model…")
-            final_pages = refine_text(raw_pages, progress)
+            progress(0.60, desc="Loading Grammar Model...")
+            final_pages = refine_grammar(raw_pages, progress)
         else:
             final_pages = raw_pages
 
-        # Step 4
-        progress(0.9, desc="Generating Word document…")
-        out_docx = os.path.join(tmp, "output.docx")
-        text_to_docx(final_pages, out_docx)
+        final_text_combined = "\n\n--- [PAGE BREAK] ---\n\n".join(final_pages)
 
-        # Copy out of the temp dir so Gradio can serve it
-        dest = tempfile.NamedTemporaryFile(
-            suffix=".docx", delete=False, prefix="scriptflow_"
-        )
-        dest.close()
-        shutil.copy(out_docx, dest.name)
+        # 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="Done!")
-    final_text = "\n\n[PAGE BREAK]\n\n".join(final_pages)
-    return raw_text, final_text, dest.name
+    progress(1.0, desc="Conversion Complete!")
+    return raw_text_combined, final_text_combined, output_file.name
 
 
 # ──────────────────────────────────────────────────────────────────────────────
-# Gradio UI
+# Gradio UI Design
 # ──────────────────────────────────────────────────────────────────────────────
 
-with gr.Blocks(title="ScriptFlow — Handwriting to Word", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="ScriptFlow — PDF to Word Converter", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
         # ✍️ ScriptFlow
-        ### Convert handwritten PDF pages into a clean Word (.docx) document
-        Upload a PDF → TrOCR reads every page → (optional) T5 corrects grammar → download .docx
+        ### 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",
+                label="Upload Handwritten PDF",
                 file_types=[".pdf"],
                 type="filepath",
             )
-            trocr_model = gr.Radio(
+            model_selector = gr.Radio(
                 choices=[
-                    ("Small — faster, ~330 MB",  "microsoft/trocr-small-handwritten"),
-                    ("Base  — more accurate, ~440 MB", "microsoft/trocr-base-handwritten"),
+                    ("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",
+                label="TrOCR Model",
             )
-            use_grammar = gr.Checkbox(
-                label="Apply T5 grammar correction (adds processing time)",
+            grammar_toggle = gr.Checkbox(
+                label="Enable T5 Grammar & Refinement",
                 value=False,
             )
-            run_btn = gr.Button("Convert", variant="primary")
+            convert_btn = gr.Button("Convert Document", variant="primary")
 
         with gr.Column(scale=2):
-            raw_out   = gr.Textbox(label="Raw OCR text",      lines=12, interactive=False)
-            final_out = gr.Textbox(label="Refined text",      lines=12, interactive=False)
-            docx_out  = gr.File(   label="Download Word file", interactive=False)
+            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(
         """
         ---
-        **Notes**
-        - First run downloads models from Hugging Face — subsequent runs use the cache.
-        - Grammar correction uses [`vennify/t5-base-grammar-correction`](https://huggingface.co/vennify/t5-base-grammar-correction).
-        - Running on CPU: expect ~30–60 s per page for OCR.
+        **Tips for best results:**
+        * High-resolution scans (200+ DPI) work significantly better.
+        * Grammar correction adds extra processing time per page on CPU.
         """
     )
 
-    run_btn.click(
-        fn=run_pipeline,
-        inputs=[pdf_input, trocr_model, use_grammar],
-        outputs=[raw_out, final_out, docx_out],
+    convert_btn.click(
+        fn=convert_pdf,
+        inputs=[pdf_input, model_selector, grammar_toggle],
+        outputs=[raw_output, refined_output, docx_download],
     )
 
 if __name__ == "__main__":