app.py

import gradio as gr
import json
import tempfile
import os
from typing import List, Optional, Literal
from PIL import Image
import spaces
from pathlib import Path
from htrflow.volume.volume import Collection
from htrflow.pipeline.pipeline import Pipeline

DEFAULT_OUTPUT = "alto"
CHOICES = ["txt", "alto", "page", "json"]

PIPELINE_CONFIGS = {
    "letter_english": {
        "steps": [
            {
                "step": "Segmentation",
                "settings": {
                    "model": "yolo",
                    "model_settings": {"model": "Riksarkivet/yolov9-lines-within-regions-1"},
                    "generation_settings": {"batch_size": 8},
                },
            },
            {
                "step": "TextRecognition",
                "settings": {
                    "model": "TrOCR",
                    "model_settings": {"model": "microsoft/trocr-base-handwritten"},
                    "generation_settings": {"batch_size": 16},
                },
            },
            {"step": "OrderLines"},
        ]
    },
    "letter_swedish": {
        "steps": [
            {
                "step": "Segmentation",
                "settings": {
                    "model": "yolo",
                    "model_settings": {"model": "Riksarkivet/yolov9-lines-within-regions-1"},
                    "generation_settings": {"batch_size": 8},
                },
            },
            {
                "step": "TextRecognition",
                "settings": {
                    "model": "TrOCR",
                    "model_settings": {"model": "Riksarkivet/trocr-base-handwritten-hist-swe-2"},
                    "generation_settings": {"batch_size": 16},
                },
            },
            {"step": "OrderLines"},
        ]
    },
    "spread_english": {
        "steps": [
            {
                "step": "Segmentation",
                "settings": {
                    "model": "yolo",
                    "model_settings": {"model": "Riksarkivet/yolov9-regions-1"},
                    "generation_settings": {"batch_size": 4},
                },
            },
            {
                "step": "Segmentation",
                "settings": {
                    "model": "yolo",
                    "model_settings": {"model": "Riksarkivet/yolov9-lines-within-regions-1"},
                    "generation_settings": {"batch_size": 8},
                },
            },
            {
                "step": "TextRecognition",
                "settings": {
                    "model": "TrOCR",
                    "model_settings": {"model": "microsoft/trocr-base-handwritten"},
                    "generation_settings": {"batch_size": 16},
                },
            },
            {"step": "ReadingOrderMarginalia", "settings": {"two_page": True}},
        ]
    },
    "spread_swedish": {
        "steps": [
            {
                "step": "Segmentation",
                "settings": {
                    "model": "yolo",
                    "model_settings": {"model": "Riksarkivet/yolov9-regions-1"},
                    "generation_settings": {"batch_size": 4},
                },
            },
            {
                "step": "Segmentation",
                "settings": {
                    "model": "yolo",
                    "model_settings": {"model": "Riksarkivet/yolov9-lines-within-regions-1"},
                    "generation_settings": {"batch_size": 8},
                },
            },
            {
                "step": "TextRecognition",
                "settings": {
                    "model": "TrOCR",
                    "model_settings": {"model": "Riksarkivet/trocr-base-handwritten-hist-swe-2"},
                    "generation_settings": {"batch_size": 16},
                },
            },
            {"step": "ReadingOrderMarginalia", "settings": {"two_page": True}},
        ]
    },
}

@spaces.GPU
def process_htr(image: Image.Image, document_type: Literal["letter_english", "letter_swedish", "spread_english", "spread_swedish"] = "letter_english", output_format: Literal["txt", "alto", "page", "json"] = DEFAULT_OUTPUT, custom_settings: Optional[str] = None):
    """Process handwritten text recognition and return extracted text with specified format file."""
    if image is None:
        return "Error: No image provided", None

    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
        image.save(temp_file.name, "PNG")
        temp_image_path = temp_file.name

    try:
        if custom_settings:
            try:
                config = json.loads(custom_settings)
            except json.JSONDecodeError:
                return "Error: Invalid JSON in custom_settings parameter", None
        else:
            config = PIPELINE_CONFIGS[document_type]

        collection = Collection([temp_image_path])
        pipeline = Pipeline.from_config(config)
        
        try:
            processed_collection = pipeline.run(collection)
        except Exception as pipeline_error:
            return f"Error: Pipeline execution failed: {str(pipeline_error)}", None

        temp_dir = Path(tempfile.mkdtemp())
        export_dir = temp_dir / output_format
        processed_collection.save(directory=str(export_dir), serializer=output_format)
        
        output_file_path = None
        for root, _, files in os.walk(export_dir):
            for file in files:
                output_file_path = os.path.join(root, file)
                break
        
        extracted_text = extract_text_from_collection(processed_collection)

        return extracted_text, output_file_path
        
    except Exception as e:
        return f"Error: HTR processing failed: {str(e)}", None
    finally:
        if os.path.exists(temp_image_path):
            os.unlink(temp_image_path)

def extract_text_from_collection(collection: Collection) -> str:
    """Extract plain text from processed collection."""
    text_lines = []
    for page in collection.pages:
        for node in page.traverse():
            if hasattr(node, "text") and node.text:
                text_lines.append(node.text)
    return "\n".join(text_lines)

def create_htrflow_mcp_server():
    demo = gr.Interface(
        fn=process_htr,
        inputs=[
            gr.Image(type="pil", label="Upload Image"),
            gr.Dropdown(choices=["letter_english", "letter_swedish", "spread_english", "spread_swedish"], value="letter_english", label="Document Type"),
            gr.Dropdown(choices=CHOICES, value=DEFAULT_OUTPUT, label="Output Format"),
            gr.Textbox(label="Custom Settings (JSON)", placeholder="Optional custom pipeline settings"),
        ],
        outputs=[
            gr.Textbox(label="Extracted Text", lines=10),
            gr.File(label="Download Output File")
        ],
        title="HTRflow MCP Server",
        description="Process handwritten text and get extracted text with output file in specified format",
        api_name="process_htr",
    )
    return demo

if __name__ == "__main__":
    demo = create_htrflow_mcp_server()
    demo.launch(mcp_server=True)