Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from kraken import blla, rpred | |
| from kraken.lib import vgsl | |
| from kraken.lib import models | |
| from PIL import Image, ImageDraw | |
| # Define available OCR models for segmentation and recognition | |
| segmentation_models = { | |
| "1col_442_sam_v1.mlmodel": "models/1col_442_sam_v1.mlmodel", | |
| "ubma_sam_v4.mlmodel": "models/ubma_sam_v4.mlmodel" | |
| } | |
| recognition_models = { | |
| "sinai_sam_rec_v4.mlmodel": "models/sinai_sam_rec_v4.mlmodel", | |
| "sinai_sam_rec_v2.mlmodel": "models/sinai_sam_rec_v2.mlmodel" | |
| } | |
| # Streamlit app title and description | |
| st.title("OCR with Kraken - Segmentation and Recognition") | |
| st.write("Upload an image, select segmentation and recognition models, and view OCR results.") | |
| # Upload image file | |
| uploaded_image = st.file_uploader("Upload an image file", type=["png", "jpg", "jpeg"]) | |
| # Select segmentation and recognition models | |
| selected_seg_model = st.selectbox("Select Kraken Segmentation Model", list(segmentation_models.keys())) | |
| selected_rec_model = st.selectbox("Select Kraken Recognition Model", list(recognition_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 models selected | |
| if uploaded_image and selected_seg_model and selected_rec_model: | |
| # Load the image | |
| image = Image.open(uploaded_image) | |
| st.image(image, caption="Uploaded Image", use_column_width=True) | |
| # Load selected Kraken segmentation and recognition models | |
| seg_model_path = segmentation_models[selected_seg_model] | |
| rec_model_path = recognition_models[selected_rec_model] | |
| seg_model = vgsl.TorchVGSLModel.load_model(seg_model_path) | |
| rec_model = models.load_any(rec_model_path) | |
| # Segment image using Kraken segmentation model | |
| baseline_seg = blla.segment(image, model=seg_model) | |
| # Pass segmentation result to recognition model | |
| pred_it = rpred.rpred(network=rec_model, im=image, bounds=baseline_seg) | |
| # Prepare to draw boundaries and display info | |
| boundaries_info = [] | |
| draw = ImageDraw.Draw(image) | |
| # Process recognition predictions for lines and draw on image | |
| for idx, pred in enumerate(pred_it): | |
| prediction = pred.prediction | |
| line_boundary = [(int(x), int(y)) for x, y in pred.boundary] | |
| line_baseline = [(int(x), int(y)) for x, y in pred.baseline] if pred.baseline else None | |
| line_type = pred.tags.get("type", "undefined") # Get line type dynamically if available | |
| # Add boundary, baseline (if selected), and prediction to display info in the new order | |
| boundaries_info.append(f"**Line {idx + 1}** (type: {line_type}):\n - Boundary: {line_boundary}") | |
| # Draw boundary in green | |
| draw.polygon(line_boundary, outline="green") | |
| # Draw baseline if the option is selected and add it to display info | |
| if draw_baselines and line_baseline: | |
| boundaries_info.append(f" - Baseline: {line_baseline}") | |
| draw.line(line_baseline, fill="red", width=2) # Draw baseline in red | |
| # Add prediction last | |
| boundaries_info.append(f" - Prediction: {prediction}") | |
| # Process and draw region boundaries from baseline_seg | |
| 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", region_type) # Get region type dynamically | |
| boundaries_info.append(f"**Region {idx + 1}** (type: {region_type_name}):\n - Boundary: {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 boundaries (green for lines, blue for regions), baselines (red if selected)", use_column_width=True) | |
| # Display the list of boundaries, predictions, and baselines | |
| st.write("**List of Boundaries, Predictions, and Baselines (if selected):**") | |
| for info in boundaries_info: | |
| st.write(info) | |