Spaces:
Sleeping
Sleeping
File size: 3,458 Bytes
528e1ce |
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 |
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)
|