johnlockejrr commited on
Commit
6fe78dc
1 Parent(s): 32f1687

Add application file

Browse files
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