|
|
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) |
