Spaces:
Running
Running
johnlockejrr
commited on
Commit
•
6fe78dc
1
Parent(s):
32f1687
Add application file
Browse files- app.py +91 -0
- models/1col_442_sam_v1.mlmodel +3 -0
- models/sinai_sam_rec_v2.mlmodel +3 -0
- models/sinai_sam_rec_v4.mlmodel +3 -0
- models/ubma_sam_v4.mlmodel +3 -0
- requirements.txt +3 -0
app.py
ADDED
@@ -0,0 +1,91 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import streamlit as st
|
2 |
+
from kraken import blla, rpred
|
3 |
+
from kraken.lib import vgsl
|
4 |
+
from kraken.lib import models
|
5 |
+
from PIL import Image, ImageDraw
|
6 |
+
|
7 |
+
# Define available OCR models for segmentation and recognition
|
8 |
+
segmentation_models = {
|
9 |
+
"1col_442_sam_v1.mlmodel": "models/1col_442_sam_v1.mlmodel",
|
10 |
+
"ubma_sam_v4.mlmodel": "models/ubma_sam_v4.mlmodel"
|
11 |
+
}
|
12 |
+
|
13 |
+
recognition_models = {
|
14 |
+
"sinai_sam_rec_v4.mlmodel": "models/sinai_sam_rec_v4.mlmodel",
|
15 |
+
"sinai_sam_rec_v2.mlmodel": "models/sinai_sam_rec_v2.mlmodel"
|
16 |
+
}
|
17 |
+
|
18 |
+
# Streamlit app title and description
|
19 |
+
st.title("OCR with Kraken - Segmentation and Recognition")
|
20 |
+
st.write("Upload an image, select segmentation and recognition models, and view OCR results.")
|
21 |
+
|
22 |
+
# Upload image file
|
23 |
+
uploaded_image = st.file_uploader("Upload an image file", type=["png", "jpg", "jpeg"])
|
24 |
+
|
25 |
+
# Select segmentation and recognition models
|
26 |
+
selected_seg_model = st.selectbox("Select Kraken Segmentation Model", list(segmentation_models.keys()))
|
27 |
+
selected_rec_model = st.selectbox("Select Kraken Recognition Model", list(recognition_models.keys()))
|
28 |
+
|
29 |
+
# Option to draw baselines
|
30 |
+
draw_baselines = st.radio("Options", ("Do not draw baselines", "Draw baselines")) == "Draw baselines"
|
31 |
+
|
32 |
+
# Process the image if uploaded and models selected
|
33 |
+
if uploaded_image and selected_seg_model and selected_rec_model:
|
34 |
+
# Load the image
|
35 |
+
image = Image.open(uploaded_image)
|
36 |
+
st.image(image, caption="Uploaded Image", use_column_width=True)
|
37 |
+
|
38 |
+
# Load selected Kraken segmentation and recognition models
|
39 |
+
seg_model_path = segmentation_models[selected_seg_model]
|
40 |
+
rec_model_path = recognition_models[selected_rec_model]
|
41 |
+
seg_model = vgsl.TorchVGSLModel.load_model(seg_model_path)
|
42 |
+
rec_model = models.load_any(rec_model_path)
|
43 |
+
|
44 |
+
# Segment image using Kraken segmentation model
|
45 |
+
baseline_seg = blla.segment(image, model=seg_model)
|
46 |
+
|
47 |
+
# Pass segmentation result to recognition model
|
48 |
+
pred_it = rpred.rpred(network=rec_model, im=image, bounds=baseline_seg)
|
49 |
+
|
50 |
+
# Prepare to draw boundaries and display info
|
51 |
+
boundaries_info = []
|
52 |
+
draw = ImageDraw.Draw(image)
|
53 |
+
|
54 |
+
# Process recognition predictions for lines and draw on image
|
55 |
+
for idx, pred in enumerate(pred_it):
|
56 |
+
prediction = pred.prediction
|
57 |
+
line_boundary = [(int(x), int(y)) for x, y in pred.boundary]
|
58 |
+
line_baseline = [(int(x), int(y)) for x, y in pred.baseline] if pred.baseline else None
|
59 |
+
line_type = pred.tags.get("type", "undefined") # Get line type dynamically if available
|
60 |
+
|
61 |
+
# Add boundary, baseline (if selected), and prediction to display info in the new order
|
62 |
+
boundaries_info.append(f"**Line {idx + 1}** (type: {line_type}):\n - Boundary: {line_boundary}")
|
63 |
+
|
64 |
+
# Draw boundary in green
|
65 |
+
draw.polygon(line_boundary, outline="green")
|
66 |
+
|
67 |
+
# Draw baseline if the option is selected and add it to display info
|
68 |
+
if draw_baselines and line_baseline:
|
69 |
+
boundaries_info.append(f" - Baseline: {line_baseline}")
|
70 |
+
draw.line(line_baseline, fill="red", width=2) # Draw baseline in red
|
71 |
+
|
72 |
+
# Add prediction last
|
73 |
+
boundaries_info.append(f" - Prediction: {prediction}")
|
74 |
+
|
75 |
+
# Process and draw region boundaries from baseline_seg
|
76 |
+
for region_type, region_list in baseline_seg.regions.items():
|
77 |
+
for idx, region_data in enumerate(region_list):
|
78 |
+
if hasattr(region_data, "boundary"):
|
79 |
+
region_boundary = [(int(x), int(y)) for x, y in region_data.boundary]
|
80 |
+
region_type_name = region_data.tags.get("type", region_type) # Get region type dynamically
|
81 |
+
boundaries_info.append(f"**Region {idx + 1}** (type: {region_type_name}):\n - Boundary: {region_boundary}")
|
82 |
+
draw.polygon(region_boundary, outline="blue") # Draw region boundary in blue
|
83 |
+
|
84 |
+
# Display the image with boundaries drawn
|
85 |
+
st.image(image, caption="Image with OCR boundaries (green for lines, blue for regions), baselines (red if selected)", use_column_width=True)
|
86 |
+
|
87 |
+
# Display the list of boundaries, predictions, and baselines
|
88 |
+
st.write("**List of Boundaries, Predictions, and Baselines (if selected):**")
|
89 |
+
for info in boundaries_info:
|
90 |
+
st.write(info)
|
91 |
+
|
models/1col_442_sam_v1.mlmodel
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:22e0b673d48337aef6ca29279418733fb861108d79ef7b3718b4c1f16eb7740e
|
3 |
+
size 5219282
|
models/sinai_sam_rec_v2.mlmodel
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:7af04abf44199e3bbb99afa2fe77f7dd6f78c4621863459b57f0e34f67a1a38c
|
3 |
+
size 16024546
|
models/sinai_sam_rec_v4.mlmodel
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:38f24382398e311d0babedab68d064b7eae8323ca9e3a2ed389102e1539926b4
|
3 |
+
size 16029821
|
models/ubma_sam_v4.mlmodel
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:900147a3c02a0b3877667eef76bba13d074b6c529013bc8aa001faf75f45e40c
|
3 |
+
size 5075650
|
requirements.txt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
kraken
|
2 |
+
pillow
|
3 |
+
streamlit
|