initial commit

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+.pth filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,25 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+ENV PIP_NO_CACHE_DIR=1
+
+COPY requirements.txt /app/requirements.txt
+
+RUN apt-get update && apt-get install -y \
+    git build-essential ffmpeg libsm6 libxext6 curl \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN pip install --upgrade pip && pip install -r requirements.txt
+
+# Detectron2 CPU build (from source) - may take time on first build
+RUN pip install 'git+https://github.com/facebookresearch/detectron2.git@v0.5#egg=detectron2'
+
+COPY app.py /app/app.py
+COPY inference.py /app/inference.py
+COPY model_final.pth /app/model_final.pth
+
+ENV USE_GPU=false
+EXPOSE 7860
+CMD ["python", "app.py"]
+
+

README.md

@@ -0,0 +1,37 @@
+## PDF OCR (Detectron2 + TrOCR) - Hugging Face Spaces
+
+This repo contains a deployable Gradio app that detects text lines with Detectron2 and reads them with TrOCR. Optional Gemini correction can refine the text.
+
+### Files
+- `app.py`: Gradio UI
+- `inference.py`: OCR pipeline (Detectron2 + TrOCR)
+- `requirements.txt`: Python dependencies (Detectron2 installed in Dockerfile)
+- `Dockerfile`: CUDA-enabled image for GPU Space
+- `model_final.pth`: Detectron2 weights
+
+### Deploy on Hugging Face Spaces (Docker Space)
+1. Create a new Space on Hugging Face → Type: Docker → Hardware: GPU (T4/A10G).
+2. Push these files to the Space repository (or connect this folder and `git push`).
+3. Set optional secret: `GEMINI_API_KEY` (for correction) in Space Settings → Secrets.
+4. Wait for the build to finish. The app will start on port 7860.
+
+### Use
+1. Upload a PDF.
+2. (Optional) Toggle Split-page (currently standard pipeline is used) and Gemini correction.
+3. Click Process.
+4. Download the ZIP of per-page JSONs. The full combined text is shown in the textbox.
+
+### Local run (GPU recommended)
+```bash
+docker build -t ocr-app .
+docker run --gpus all -p 7860:7860 ocr-app
+```
+
+Then open http://localhost:7860
+
+### Notes
+- Detectron2 requires GPU for reasonable speed; CPU will be slow.
+- `TEXTLINE_MODEL_PATH` can be overridden via env var if the weights are elsewhere.
+- TrOCR models are downloaded on first run and cached in the container layer after warmup.
+
+

app.py

@@ -0,0 +1,32 @@
+import os
+import gradio as gr
+from inference import run_ocr
+
+
+def predict(pdf_file, split_page, use_llm, gemini_key):
+    if pdf_file is None:
+        return "Please upload a PDF.", None
+    key = gemini_key or os.getenv("GEMINI_API_KEY", None)
+    text, zip_path = run_ocr(pdf_file.name, split_page_enabled=split_page, use_llm=use_llm, gemini_key=key)
+    return text, zip_path
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("## PDF OCR (Detectron2 + TrOCR)")
+    with gr.Row():
+        with gr.Column():
+            pdf = gr.File(label="Upload PDF", file_types=[".pdf"])
+            split_page = gr.Checkbox(label="Split-page mode", value=False)
+            use_llm = gr.Checkbox(label="Gemini correction", value=False)
+            gemini_key = gr.Textbox(label="Gemini API Key (optional)", type="password")
+            btn = gr.Button("Process")
+        with gr.Column():
+            out_text = gr.Textbox(label="Extracted Text", lines=18)
+            out_zip = gr.File(label="Per-page JSON (ZIP)")
+    btn.click(predict, inputs=[pdf, split_page, use_llm, gemini_key], outputs=[out_text, out_zip])
+
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)
+
+

inference.py

@@ -0,0 +1,201 @@
+import os
+import io
+import json
+import time
+import shutil
+import tempfile
+from typing import Tuple
+
+import cv2
+import fitz  # PyMuPDF
+import numpy as np
+from PIL import Image
+
+import torch
+from detectron2.config import get_cfg
+from detectron2.engine import DefaultPredictor
+from detectron2.data import MetadataCatalog
+from detectron2 import model_zoo
+
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+
+
+# -----------------------------
+# Configuration (override via env if needed)
+# -----------------------------
+TEXTLINE_MODEL_PATH = os.getenv("TEXTLINE_MODEL_PATH", "./model_final.pth")
+USE_GPU = os.getenv("USE_GPU", "true").lower() == "true"
+SCORE_THRESHOLD = float(os.getenv("SCORE_THRESHOLD", "0.5"))
+AREA_THRESHOLD_PERCENT = float(os.getenv("AREA_THRESHOLD_PERCENT", "12.5"))
+DPI = int(os.getenv("PDF_DPI", "200"))
+
+TROCR_SPANISH_MODEL = os.getenv("TROCR_SPANISH_MODEL", "qantev/trocr-large-spanish")
+TROCR_FALLBACK_MODEL = os.getenv("TROCR_FALLBACK_MODEL", "microsoft/trocr-base-printed")
+
+
+class EnhancedTextlineExtractor:
+    def __init__(self, model_path: str):
+        self.cfg = self._setup_cfg(model_path)
+        self.predictor = DefaultPredictor(self.cfg)
+
+        # Init TrOCR
+        self.device = torch.device("cuda" if torch.cuda.is_available() and USE_GPU else "cpu")
+        self.trocr_processor, self.trocr_model = self._load_trocr()
+        self.trocr_model.to(self.device)
+
+    def _setup_cfg(self, model_path: str):
+        cfg = get_cfg()
+        cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml"))
+        cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2  # textline, baseline
+        cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = SCORE_THRESHOLD
+        cfg.MODEL.WEIGHTS = model_path
+        cfg.DATASETS.TEST = ("page_test",)
+        cfg.DATALOADER.NUM_WORKERS = 2
+        MetadataCatalog.get("page_test").thing_classes = ["textline", "baseline"]
+        return cfg
+
+    def _load_trocr(self):
+        try:
+            processor = TrOCRProcessor.from_pretrained(TROCR_SPANISH_MODEL)
+            model = VisionEncoderDecoderModel.from_pretrained(TROCR_SPANISH_MODEL)
+            return processor, model
+        except Exception:
+            processor = TrOCRProcessor.from_pretrained(TROCR_FALLBACK_MODEL)
+            model = VisionEncoderDecoderModel.from_pretrained(TROCR_FALLBACK_MODEL)
+            return processor, model
+
+    def pdf_to_images(self, pdf_path: str, dpi: int = DPI):
+        doc = fitz.open(pdf_path)
+        images = []
+        try:
+            for page_num in range(len(doc)):
+                page = doc.load_page(page_num)
+                mat = fitz.Matrix(dpi / 72, dpi / 72)
+                pix = page.get_pixmap(matrix=mat)
+                img_data = pix.tobytes("png")
+                nparr = np.frombuffer(img_data, np.uint8)
+                img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+                images.append(img)
+        finally:
+            doc.close()
+        return images
+
+    def filter_margin_boxes_by_area(self, boxes, scores, area_threshold_percent: float = AREA_THRESHOLD_PERCENT):
+        if len(boxes) == 0:
+            return np.array([]), np.array([]), np.array([]), np.array([])
+        areas = []
+        for box in boxes:
+            x1, y1, x2, y2 = box
+            areas.append((x2 - x1) * (y2 - y1))
+        areas = np.array(areas)
+        avg_area = np.mean(areas)
+        area_threshold = avg_area * (area_threshold_percent / 100.0)
+        main_boxes, main_scores, margin_boxes, margin_scores = [], [], [], []
+        for b, s, a in zip(boxes, scores, areas):
+            if a >= area_threshold:
+                main_boxes.append(b)
+                main_scores.append(s)
+            else:
+                margin_boxes.append(b)
+                margin_scores.append(s)
+        return np.array(main_boxes), np.array(main_scores), np.array(margin_boxes), np.array(margin_scores)
+
+    def process_page_standard(self, image):
+        outputs = self.predictor(image)
+        instances = outputs["instances"]
+        boxes = instances.pred_boxes.tensor.cpu().numpy()
+        scores = instances.scores.cpu().numpy()
+        if len(boxes) == 0:
+            return {"success": False, "error": "No textlines detected"}
+        main_boxes, main_scores, _, _ = self.filter_margin_boxes_by_area(boxes, scores)
+        if len(main_boxes) == 0:
+            return {"success": False, "error": "No textlines after filtering"}
+
+        line_segments = []
+        full_text_lines = []
+        for i, (box, score) in enumerate(zip(main_boxes, main_scores)):
+            x1, y1, x2, y2 = map(int, box)
+            crop_bgr = image[y1:y2, x1:x2]
+            try:
+                crop_rgb = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2RGB)
+                pil_image = Image.fromarray(crop_rgb)
+                pixel_values = self.trocr_processor(images=pil_image, return_tensors="pt").pixel_values
+                pixel_values = pixel_values.to(self.device)
+                with torch.no_grad():
+                    generated_ids = self.trocr_model.generate(pixel_values, max_new_tokens=128)
+                    generated_text = self.trocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+                text = generated_text.strip()
+                full_text_lines.append(text)
+                line_segments.append({
+                    "line_index": i,
+                    "bbox": [int(x1), int(y1), int(x2), int(y2)],
+                    "score": float(score),
+                    "text": text,
+                    "confidence": 1.0
+                })
+            except Exception:
+                line_segments.append({
+                    "line_index": i,
+                    "bbox": [int(x1), int(y1), int(x2), int(y2)],
+                    "score": float(score),
+                    "text": "",
+                    "confidence": 0.0
+                })
+        return {
+            "success": True,
+            "line_segments": line_segments,
+            "full_text": "\n".join(full_text_lines)
+        }
+
+
+def _zip_directory(src_dir: str, zip_path: str) -> str:
+    base, _ = os.path.splitext(zip_path)
+    archive = shutil.make_archive(base, 'zip', src_dir)
+    return archive
+
+
+def run_ocr(pdf_path: str, split_page_enabled: bool = False, use_llm: bool = False, gemini_key: str = None) -> Tuple[str, str]:
+    """
+    Run OCR on the provided PDF.
+
+    Returns:
+      combined_text (str), zip_file_path (str)
+    """
+    extractor = EnhancedTextlineExtractor(TEXTLINE_MODEL_PATH)
+    images = extractor.pdf_to_images(pdf_path, dpi=DPI)
+
+    temp_dir = tempfile.mkdtemp(prefix="ocr_outputs_")
+    inferences_dir = os.path.join(temp_dir, "inferences")
+    os.makedirs(inferences_dir, exist_ok=True)
+
+    all_results = []
+    for i, image in enumerate(images):
+        result = extractor.process_page_standard(image)
+        all_results.append(result)
+        page_file = os.path.join(inferences_dir, f"page_{i+1}_result.json")
+        with open(page_file, "w", encoding="utf-8") as f:
+            json.dump(result, f, ensure_ascii=False, indent=2)
+
+    combined_text = "\n\n".join([r.get("full_text", "") for r in all_results if r.get("success")])
+
+    # Optional Gemini correction over combined text (simple, single pass)
+    if use_llm and gemini_key and combined_text.strip():
+        try:
+            import google.generativeai as genai
+            genai.configure(api_key=gemini_key)
+            prompt = (
+                "Correct the following historical Spanish OCR text while preserving grammar and style. "
+                "Fix orthography, punctuation, and obvious OCR mistakes. Return only corrected text.\n\n" + combined_text
+            )
+            response = genai.GenerativeModel('gemini-2.5-pro').generate_content(prompt)
+            if getattr(response, 'text', None):
+                combined_text = response.text.strip()
+        except Exception:
+            # Swallow LLM errors and return original text
+            pass
+
+    zip_path = os.path.join(temp_dir, "per_page_jsons.zip")
+    archive_path = _zip_directory(inferences_dir, zip_path)
+    return combined_text, archive_path
+
+

model_final.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:466b008d261466deff5ab6f0517403441bcf7e379f39a715709b78503d252158
+size 503106240

requirements.txt

@@ -0,0 +1,16 @@
+torch==2.6.0
+torchvision==0.21.0
+transformers==4.44.2
+opencv-python-headless==4.11.0.86
+pillow==11.2.1
+layoutparser==0.3.4
+pdfplumber==0.11.7
+pymupdf==1.24.10
+numpy==1.26.4
+scipy==1.15.3
+pandas==2.2.3
+google-generativeai==0.7.2
+python-dotenv==1.0.1
+gradio==4.44.0
+# detectron2 installed via Dockerfile to match CUDA
+