import streamlit as st
from kraken import blla
from kraken.lib import vgsl
from PIL import Image, ImageDraw

# Define available OCR models
ocr_models = {
        "blla.mlmodel": "models/blla.mlmodel",
        "syrnt_blla_v1_best.mlmodel": "models/syrnt_blla_v1_best.mlmodel",
        "syrnt_blla_v2_best.mlmodel": "models/syrnt_blla_v2_best.mlmodel",
        "syrnt_cl_v1_best.mlmodel": "models/syrnt_cl_v1_best.mlmodel",
        "syrnt_cl_v2_best.mlmodel": "models/syrnt_cl_v2_best.mlmodel",
        "syrnt_two_blla_v1_best.mlmodel": "models/syrnt_two_blla_v1_best.mlmodel",
        "syrnt_two_v1_best.mlmodel": "models/syrnt_two_v1_best.mlmodel"
}

# Streamlit app title and description
st.title("OCR with Kraken")
st.write("Upload an image, select an OCR model, and view detected polygons.")

# Upload image file
uploaded_image = st.file_uploader("Upload an image file", type=["png", "jpg", "jpeg"])

# Select model from dropdown
selected_model = st.selectbox("Select Kraken OCR model", list(ocr_models.keys()))

# Option to draw baselines
draw_baselines = st.radio("Options", ("Do not draw baselines", "Draw baselines")) == "Draw baselines"

# Process the image if uploaded and model selected
if uploaded_image and selected_model:
    # Load the image
    image = Image.open(uploaded_image)
    st.image(image, caption="Uploaded Image", use_column_width=True)
    
    # Load selected Kraken model
    model_path = ocr_models[selected_model]
    model = vgsl.TorchVGSLModel.load_model(model_path)

    # Segment image using Kraken
    baseline_seg = blla.segment(image, model=model)

    # Prepare to draw boundaries and display info
    boundaries_info = []
    draw = ImageDraw.Draw(image)

    # Process and draw line boundaries
    for idx, line_data in enumerate(baseline_seg.lines):
        line_boundary = [(int(x), int(y)) for x, y in line_data.boundary]
        line_type = line_data.tags.get("type", "undefined")  # Dynamically get line type name
        boundaries_info.append(f"Line {idx + 1} (type: {line_type}): {line_boundary}")
        draw.polygon(line_boundary, outline="green")  # Draw line boundary in green

        # Draw baseline if the option is selected
        if draw_baselines:
            line_baseline = [(int(x), int(y)) for x, y in line_data.baseline]
            boundaries_info.append(f"  Baseline for Line {idx + 1}: {line_baseline}")
            draw.line(line_baseline, fill="red", width=2)  # Draw baseline in red

    # Process and draw region boundaries by iterating through each region type
    for region_type, region_list in baseline_seg.regions.items():
        for idx, region_data in enumerate(region_list):
            if hasattr(region_data, "boundary"):
                region_boundary = [(int(x), int(y)) for x, y in region_data.boundary]
                region_type_name = region_data.tags.get("type", "undefined")  # Get region type dynamically
                boundaries_info.append(f"Region {idx + 1} (type: {region_type_name}): {region_boundary}")
                draw.polygon(region_boundary, outline="blue")  # Draw region boundary in blue

    # Display the image with boundaries drawn
    st.image(image, caption="Image with OCR lines (green), regions (blue), and baselines (red if selected)", use_column_width=True)

    # Display the list of boundaries with line and region types
    st.write("List of Boundaries (Lines and Regions):")
    for info in boundaries_info:
        st.write(info)