app.py

"""
Kiri OCR - Gradio Demo for Hugging Face Spaces

A lightweight OCR library for English and Khmer documents with streaming output support.
"""
import gradio as gr
import numpy as np
from PIL import Image
import cv2
import tempfile
import os

# Global OCR instances (one per decode method)
ocr_instances = {}

def load_ocr(decode_method="accurate"):
    """Load the OCR model with specified decode method."""
    from kiri_ocr import OCR
    print(f"Loading OCR model with decode_method={decode_method}...")
    return OCR(
        model_path="mrrtmob/kiri-ocr",
        det_method="db",
        decode_method=decode_method,
        device="cpu",
        verbose=False
    )

def get_ocr(decode_method="accurate"):
    """Get or create OCR instance for the specified decode method."""
    global ocr_instances
    if decode_method not in ocr_instances:
        ocr_instances[decode_method] = load_ocr(decode_method)
    return ocr_instances[decode_method]

def process_document_stream(image, decode_method):
    """
    Process document image with real-time character streaming.
    
    Args:
        image: Input image (PIL Image or numpy array)
        decode_method: Decode method to use (fast, accurate, or beam)
    
    Yields:
        Tuple of (annotated_image, extracted_text)
    """
    if image is None:
        yield None, "Please upload an image."
        return
        
    try:
        ocr_engine = get_ocr(decode_method)
        
        # Save temp file for processing (required by current API)
        # Convert PIL to BGR numpy array first if needed
        if isinstance(image, Image.Image):
            img_array = np.array(image)
        else:
            img_array = image
            
        # Handle channels
        if len(img_array.shape) == 2:
            img_display = cv2.cvtColor(img_array, cv2.COLOR_GRAY2BGR)
        elif img_array.shape[2] == 4:
            img_display = cv2.cvtColor(img_array, cv2.COLOR_RGBA2BGR)
        else:
            img_display = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
            
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
            temp_path = f.name
            
        cv2.imwrite(temp_path, img_display)
        
        # Variables for state tracking
        annotated = img_display.copy()
        extracted_text = ""
        current_region_text = ""
        
        # Use the streaming generator
        for chunk in ocr_engine.extract_text_stream_chars(temp_path, mode="lines"):
            
            # Handle region boundaries
            if chunk.get("region_start"):
                # Draw box for new region
                if "box" in chunk:
                    x, y, w, h = chunk["box"]
                    # Draw box
                    cv2.rectangle(annotated, (x, y), (x + w, y + h), (0, 255, 0), 2)
                    # Draw region number
                    cv2.putText(
                        annotated, str(chunk["region_number"]), (x, y - 5),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1
                    )
                
                # Add newline if not first region
                if chunk["region_number"] > 1:
                    extracted_text += "\n"
            
            # Append new token
            token = chunk.get("token", "")
            if token:
                extracted_text += token
                current_region_text += token
            
            # Update display every few chars or at region boundaries to keep UI responsive
            # (Gradio streaming works best with frequent updates)
            if chunk.get("region_start") or chunk.get("region_finished") or len(current_region_text) % 3 == 0:
                # Convert BGR back to RGB for Gradio
                yield cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB), extracted_text
                
        # Final update
        yield cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB), extracted_text
        
        # Cleanup
        os.unlink(temp_path)
        
    except Exception as e:
        import traceback
        yield image, f"Error: {str(e)}\n{traceback.format_exc()}"


def recognize_line_stream(image, decode_method):
    """
    Stream text from single line image.
    
    Args:
        image: Input image
        decode_method: Decode method to use (fast, accurate, or beam)
    """
    if image is None:
        yield "Please upload an image."
        return
        
    try:
        ocr_engine = get_ocr(decode_method)
        
        # Save temp file
        if isinstance(image, Image.Image):
            image.save("temp_line.png")
            path = "temp_line.png"
        else:
            cv2.imwrite("temp_line.png", cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
            path = "temp_line.png"
            
        extracted_text = ""
        
        for chunk in ocr_engine.recognize_streaming(path):
            token = chunk.get("token", "")
            if token:
                extracted_text += token
                yield extracted_text
                
        if os.path.exists(path):
            os.unlink(path)
            
    except Exception as e:
        yield f"Error: {str(e)}"

# Custom CSS
css = """
.container { max-width: 1200px; margin: auto; }
.output-text { font-family: monospace; }
"""

# Create Gradio interface
with gr.Blocks(title="Kiri OCR - Streaming Demo", css=css, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # ⚡ Kiri OCR Demo
        
        **Real-time OCR for English and Khmer documents**
        
        This demo showcases the **character-by-character streaming** capability of Kiri OCR.
        """
    )
    
    with gr.Tabs():
        # Document OCR Tab
        with gr.TabItem("📄 Document Stream"):
            gr.Markdown("Upload a document to see text appear in real-time as it's recognized.")
            
            with gr.Row():
                with gr.Column(scale=1):
                    doc_input = gr.Image(
                        label="Upload Document",
                        type="pil",
                        sources=["upload", "clipboard", "webcam"]
                    )
                    
                    # Decode method selector
                    doc_decode_method = gr.Radio(
                        choices=["fast", "accurate", "beam"],
                        value="accurate",
                        label="Decode Method",
                        info="Fast: Fastest, lower accuracy | Accurate: Balanced | Beam: Slowest, highest accuracy"
                    )
                    
                    with gr.Row():
                        doc_btn = gr.Button("⚡ Stream Text", variant="primary")
                        doc_stop = gr.Button("⏹️ Stop", variant="secondary", visible=False)
                
                with gr.Column(scale=1):
                    # Annotated image updates in real-time
                    doc_output_img = gr.Image(label="Live Detection")
                    # Text updates character-by-character
                    doc_output_text = gr.Textbox(
                        label="Streaming Text",
                        lines=15,
                        autoscroll=True,
                        elem_classes=["output-text"]
                    )
            
        # Single Line OCR Tab
        with gr.TabItem("✏️ Single Line Stream"):
            gr.Markdown("Stream text recognition for a single cropped text line.")
            
            with gr.Row():
                with gr.Column(scale=1):
                    line_input = gr.Image(
                        label="Upload Text Line",
                        type="pil",
                        sources=["upload", "clipboard"]
                    )
                    
                    # Decode method selector
                    line_decode_method = gr.Radio(
                        choices=["fast", "accurate", "beam"],
                        value="accurate",
                        label="Decode Method",
                        info="Fast: Fastest, lower accuracy | Accurate: Balanced | Beam: Slowest, highest accuracy"
                    )
                    
                    with gr.Row():
                        line_btn = gr.Button("⚡ Stream Recognize", variant="primary")
                        line_stop = gr.Button("⏹️ Stop", variant="secondary", visible=False)
                
                with gr.Column(scale=1):
                    line_output_text = gr.Textbox(
                        label="Streaming Output",
                        lines=3,
                        elem_classes=["output-text"]
                    )
            
    
    # Toggle buttons visibility
    def toggle_doc_buttons():
        return gr.update(visible=False), gr.update(visible=True)
    
    def reset_doc_buttons():
        return gr.update(visible=True), gr.update(visible=False)

    doc_event = doc_btn.click(
        fn=toggle_doc_buttons,
        outputs=[doc_btn, doc_stop]
    ).then(
        fn=process_document_stream,
        inputs=[doc_input, doc_decode_method],
        outputs=[doc_output_img, doc_output_text]
    ).then(
        fn=reset_doc_buttons,
        outputs=[doc_btn, doc_stop]
    )
    
    doc_stop.click(
        fn=reset_doc_buttons,
        outputs=[doc_btn, doc_stop],
        cancels=[doc_event]
    )
    
    def toggle_line_buttons():
        return gr.update(visible=False), gr.update(visible=True)
        
    def reset_line_buttons():
        return gr.update(visible=True), gr.update(visible=False)

    line_event = line_btn.click(
        fn=toggle_line_buttons,
        outputs=[line_btn, line_stop]
    ).then(
        fn=recognize_line_stream,
        inputs=[line_input, line_decode_method],
        outputs=line_output_text
    ).then(
        fn=reset_line_buttons,
        outputs=[line_btn, line_stop]
    )
    
    line_stop.click(
        fn=reset_line_buttons,
        outputs=[line_btn, line_stop],
        cancels=[line_event]
    )

    gr.Markdown(
        """
        ---
        ### 🔍 Decode Methods:
        - **Fast**: Greedy decoding - fastest speed, good for quick previews
        - **Accurate**: Default balanced mode - good speed and accuracy
        - **Beam**: Beam search decoding - slowest but highest accuracy
        
        ---
        [GitHub Repository](https://github.com/mrrtmob/kiri-ocr) | [Hugging Face Model](https://huggingface.co/mrrtmob/kiri-ocr)
        """
    )

# Launch
if __name__ == "__main__":
    demo.queue().launch()