Spaces:
Sleeping
Sleeping
File size: 4,199 Bytes
6fe78dc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
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)
|