import gradio as gr import json import tempfile import os from typing import List, Optional, Literal, Tuple 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" FORMAT_CHOICES = ["letter_english", "letter_swedish", "spread_english", "spread_swedish"] FILE_CHOICES = ["txt", "alto", "page", "json"] FormatChoices = Literal["letter_english", "letter_swedish", "spread_english", "spread_swedish"] FileChoices = Literal["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 htrflow_htr_url(image_path: str, document_type: FormatChoices = "letter_swedish", output_format: FileChoices = DEFAULT_OUTPUT, custom_settings: Optional[str] = None, server_name: str = "https://gabriel-htrflow-mcp.hf.space") -> Tuple[str, str]: """ Process handwritten text recognition (HTR) on uploaded images and return both file content and download link. This function uses machine learning models to automatically detect, segment, and transcribe handwritten text from historical documents. It supports different document types and languages, with specialized models trained on historical handwriting from the Swedish National Archives (Riksarkivet). Args: image_path (str): The file path or URL to the image containing handwritten text to be processed. Supports common image formats like JPG, PNG, TIFF. document_type (FormatChoices): The type of document and language processing template to use. Available options: - "letter_english": Single-page English handwritten letters - "letter_swedish": Single-page Swedish handwritten letters (default) - "spread_english": Two-page spread English documents with marginalia - "spread_swedish": Two-page spread Swedish documents with marginalia Default: "letter_swedish" output_format (FileChoices): The format for the output file containing the transcribed text. Available options: - "txt": Plain text format with line breaks - "alto": ALTO XML format with detailed layout and coordinate information - "page": PAGE XML format with structural markup and positioning data - "json": JSON format with structured text, layout information and metadata Default: "alto" custom_settings (Optional[str]): Advanced users can provide custom pipeline configuration as a JSON string to override the default processing steps. Default: None (uses predefined configuration for document_type) server_name (str): The base URL of the server for constructing download links. Default: "https://gabriel-htrflow-mcp.hf.space" Returns: Tuple[str, str]: A tuple containing: - JSON string with extracted text, file content - File path for direct download via gr.File (server_name/gradio_api/file=/tmp/gradio/{temp_folder}/{file_name}) """ if not image_path: error_json = json.dumps({"error": "No image provided"}) return error_json, None try: original_filename = Path(image_path).stem or "output" if custom_settings: try: config = json.loads(custom_settings) except json.JSONDecodeError: error_json = json.dumps({"error": "Invalid JSON in custom_settings parameter"}) return error_json, None else: config = PIPELINE_CONFIGS[document_type] collection = Collection([image_path]) pipeline = Pipeline.from_config(config) try: processed_collection = pipeline.run(collection) except Exception as pipeline_error: error_json = json.dumps({"error": f"Pipeline execution failed: {str(pipeline_error)}"}) return error_json, None extracted_text = extract_text_from_collection(processed_collection) 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: old_path = os.path.join(root, file) file_ext = Path(file).suffix new_filename = f"{original_filename}.{output_format}" if not file_ext else f"{original_filename}{file_ext}" new_path = os.path.join(root, new_filename) os.rename(old_path, new_path) output_file_path = new_path break if output_file_path and os.path.exists(output_file_path): with open(output_file_path, 'r', encoding='utf-8') as f: file_content = f.read() result = { "text": extracted_text, "content": file_content, } json_result = json.dumps(result, ensure_ascii=False, indent=2) return json_result, output_file_path else: error_json = json.dumps({"error": "Failed to generate output file"}) return error_json, None except Exception as e: error_json = json.dumps({"error": f"HTR processing failed: {str(e)}"}) return error_json, None def htrflow_visualizer(image: str, htr_document: str) -> str: pass def extract_text_from_collection(collection: Collection) -> str: 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(): htrflow_url = gr.Interface( fn=htrflow_htr_url, inputs=[ gr.Image(type="filepath", label="Upload Image or Enter URL"), gr.Dropdown(choices=FORMAT_CHOICES, value="letter_swedish", label="Document Type"), gr.Dropdown(choices=FILE_CHOICES, value=DEFAULT_OUTPUT, label="Output Format"), gr.Textbox(label="Custom Settings (JSON)", placeholder="Optional custom pipeline settings", value=""), ], outputs=[ gr.Textbox(label="HTR Result (JSON)", lines=10), gr.File(label="Download HTR Output File") ], description="Process handwritten text from uploaded file or URL and get both content and download link for file", api_name="htrflow_htr_url", ) htrflow_viz = gr.Interface( fn=htrflow_visualizer, inputs=[ gr.Image(type="filepath", label="Upload Image or Enter URL"), gr.Textbox(label="HTR Document content", placeholder="Path to the HTR document file", value=""), ], outputs=gr.File(label="Download Output File"), description="Visualize document", api_name="htrflow_visualizer" ) demo = gr.TabbedInterface( [htrflow_url, htrflow_viz], ["HTR URL", "HTR Visualizer"], title="HTRflow Handwritten Text Recognition", ) return demo if __name__ == "__main__": demo = create_htrflow_mcp_server() demo.launch(mcp_server=True, share=False, debug=False)