app.py

import gradio as gr
import json
import base64
import tempfile
import os
from typing import Dict, List, Optional, Literal
from datetime import datetime
from PIL import Image, ImageDraw, ImageFont
import io
import spaces
import shutil
from pathlib import Path
from htrflow.volume.volume import Collection
from htrflow.pipeline.pipeline import Pipeline

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"] = "spread_swedish", confidence_threshold: float = 0.8, custom_settings: Optional[str] = None) -> Dict:
    """Process handwritten text recognition on uploaded images using HTRflow pipelines."""
    try:
        if image is None:
            return {"success": False, "error": "No image provided", "results": 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 {"success": False, "error": "Invalid JSON in custom_settings parameter", "results": None}
            else:
                config = PIPELINE_CONFIGS[document_type]

            collection = Collection([temp_image_path])
            pipeline = Pipeline.from_config(config)
            processed_collection = pipeline.run(collection)

            img_buffer = io.BytesIO()
            image.save(img_buffer, format="PNG")
            image_base64 = base64.b64encode(img_buffer.getvalue()).decode("utf-8")

            results = extract_text_results(processed_collection, confidence_threshold)
            
            processing_state = {
                "collection_data": serialize_collection_data(processed_collection),
                "image_base64": image_base64,
                "image_size": image.size,
                "document_type": document_type,
                "confidence_threshold": confidence_threshold,
                "timestamp": datetime.now().isoformat(),
            }

            return {
                "success": True,
                "results": results,
                "processing_state": json.dumps(processing_state),
                "metadata": {
                    "total_lines": len(results.get("text_lines", [])),
                    "average_confidence": results.get("average_confidence", 0),
                    "document_type": document_type,
                    "image_dimensions": image.size,
                },
            }
        finally:
            if os.path.exists(temp_image_path):
                os.unlink(temp_image_path)
    except Exception as e:
        return {"success": False, "error": f"HTR processing failed: {str(e)}", "results": None}

def visualize_results(processing_state: str, visualization_type: Literal["overlay", "confidence_heatmap", "text_regions"] = "overlay", show_confidence: bool = True, highlight_low_confidence: bool = True, image: Optional[Image.Image] = None) -> Dict:
    """Generate interactive visualizations of HTR processing results."""
    try:
        state = json.loads(processing_state)
        collection_data = state["collection_data"]

        if image is not None:
            original_image = image
        else:
            image_data = base64.b64decode(state["image_base64"])
            original_image = Image.open(io.BytesIO(image_data))

        viz_image = create_visualization(original_image, collection_data, visualization_type, show_confidence, highlight_low_confidence)

        img_buffer = io.BytesIO()
        viz_image.save(img_buffer, format="PNG")
        img_base64 = base64.b64encode(img_buffer.getvalue()).decode("utf-8")

        return {
            "success": True,
            "visualization": {
                "image_base64": img_base64,
                "image_format": "PNG",
                "visualization_type": visualization_type,
                "dimensions": viz_image.size,
            },
            "metadata": {"total_elements": len(collection_data.get("text_elements", []))},
        }
    except Exception as e:
        return {"success": False, "error": f"Visualization generation failed: {str(e)}", "visualization": None}

def export_results(processing_state: str, output_formats: List[Literal["txt", "json", "alto", "page"]] = ["txt"], confidence_filter: float = 0.0) -> Dict:
    """Export HTR results to multiple formats using HTRflow's native export functionality."""
    try:
        state = json.loads(processing_state)
        
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
            image_data = base64.b64decode(state["image_base64"])
            image = Image.open(io.BytesIO(image_data))
            image.save(temp_file.name, "PNG")
            temp_image_path = temp_file.name

        try:
            collection = Collection([temp_image_path])
            pipeline = Pipeline.from_config(PIPELINE_CONFIGS[state["document_type"]])
            processed_collection = pipeline.run(collection)

            temp_dir = Path(tempfile.mkdtemp())
            exports = {}
            
            for fmt in output_formats:
                export_dir = temp_dir / fmt
                processed_collection.save(directory=str(export_dir), serializer=fmt)
                
                export_files = []
                for root, _, files in os.walk(export_dir):
                    for file in files:
                        file_path = os.path.join(root, file)
                        with open(file_path, 'r', encoding='utf-8') as f:
                            content = f.read()
                        export_files.append({"filename": file, "content": content})
                
                exports[fmt] = export_files

            shutil.rmtree(temp_dir)

            return {
                "success": True,
                "exports": exports,
                "export_metadata": {
                    "formats_generated": output_formats,
                    "confidence_filter": confidence_filter,
                    "timestamp": datetime.now().isoformat(),
                },
            }
        finally:
            if os.path.exists(temp_image_path):
                os.unlink(temp_image_path)
                
    except Exception as e:
        return {"success": False, "error": f"Export generation failed: {str(e)}", "exports": None}

def extract_text_results(collection: Collection, confidence_threshold: float) -> Dict:
    results = {"extracted_text": "", "text_lines": [], "confidence_scores": []}
    for page in collection.pages:
        for node in page.traverse():
            if hasattr(node, "text") and node.text and hasattr(node, "confidence") and node.confidence >= confidence_threshold:
                results["text_lines"].append({
                    "text": node.text,
                    "confidence": node.confidence,
                    "bbox": getattr(node, "bbox", None),
                })
                results["extracted_text"] += node.text + "\n"
                results["confidence_scores"].append(node.confidence)
    
    results["average_confidence"] = sum(results["confidence_scores"]) / len(results["confidence_scores"]) if results["confidence_scores"] else 0
    return results

def serialize_collection_data(collection: Collection) -> Dict:
    text_elements = []
    for page in collection.pages:
        for node in page.traverse():
            if hasattr(node, "text") and node.text:
                text_elements.append({
                    "text": node.text,
                    "confidence": getattr(node, "confidence", 1.0),
                    "bbox": getattr(node, "bbox", None),
                })
    return {"text_elements": text_elements}

def create_visualization(image, collection_data, visualization_type, show_confidence, highlight_low_confidence):
    viz_image = image.copy()
    draw = ImageDraw.Draw(viz_image)
    
    try:
        font = ImageFont.truetype("arial.ttf", 12)
    except:
        font = ImageFont.load_default()

    for element in collection_data.get("text_elements", []):
        if element.get("bbox"):
            bbox = element["bbox"]
            confidence = element.get("confidence", 1.0)
            
            if visualization_type == "overlay":
                color = (255, 165, 0) if highlight_low_confidence and confidence < 0.7 else (0, 255, 0)
                draw.rectangle(bbox, outline=color, width=2)
                if show_confidence:
                    draw.text((bbox[0], bbox[1] - 15), f"{confidence:.2f}", fill=color, font=font)
            
            elif visualization_type == "confidence_heatmap":
                if confidence < 0.5:
                    color = (255, 0, 0, 100)
                elif confidence < 0.8:
                    color = (255, 255, 0, 100)
                else:
                    color = (0, 255, 0, 100)
                overlay = Image.new("RGBA", viz_image.size, (0, 0, 0, 0))
                overlay_draw = ImageDraw.Draw(overlay)
                overlay_draw.rectangle(bbox, fill=color)
                viz_image = Image.alpha_composite(viz_image.convert("RGBA"), overlay)
            
            elif visualization_type == "text_regions":
                colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0)]
                color = colors[hash(str(bbox)) % len(colors)]
                draw.rectangle(bbox, outline=color, width=3)

    return viz_image.convert("RGB") if visualization_type == "confidence_heatmap" else viz_image

def create_htrflow_mcp_server():
    demo = gr.TabbedInterface(
        [
            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.Slider(0.0, 1.0, value=0.8, label="Confidence Threshold"),
                    gr.Textbox(label="Custom Settings (JSON)", placeholder="Optional custom pipeline settings"),
                ],
                outputs=gr.JSON(label="Processing Results"),
                title="HTR Processing Tool",
                description="Process handwritten text using configurable HTRflow pipelines",
                api_name="process_htr",
            ),
            gr.Interface(
                fn=visualize_results,
                inputs=[
                    gr.Textbox(label="Processing State (JSON)", placeholder="Paste processing results from HTR tool"),
                    gr.Dropdown(choices=["overlay", "confidence_heatmap", "text_regions"], value="overlay", label="Visualization Type"),
                    gr.Checkbox(value=True, label="Show Confidence Scores"),
                    gr.Checkbox(value=True, label="Highlight Low Confidence"),
                    gr.Image(type="pil", label="Image (optional)"),
                ],
                outputs=gr.JSON(label="Visualization Results"),
                title="Results Visualization Tool",
                description="Generate interactive visualizations of HTR results",
                api_name="visualize_results",
            ),
            gr.Interface(
                fn=export_results,
                inputs=[
                    gr.Textbox(label="Processing State (JSON)", placeholder="Paste processing results from HTR tool"),
                    gr.CheckboxGroup(choices=["txt", "json", "alto", "page"], value=["txt"], label="Output Formats"),
                    gr.Slider(0.0, 1.0, value=0.0, label="Confidence Filter"),
                ],
                outputs=gr.JSON(label="Export Results"),
                title="Export Tool",
                description="Export HTR results to multiple formats",
                api_name="export_results",
            ),
        ],
        ["HTR Processing", "Results Visualization", "Export Results"],
        title="HTRflow MCP Server",
    )
    return demo

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