Spaces:
Sleeping
Sleeping
| 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) | |