-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathstreamlit_app.py
More file actions
168 lines (147 loc) · 7 KB
/
streamlit_app.py
File metadata and controls
168 lines (147 loc) · 7 KB
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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
import streamlit as st
import cv2
import matplotlib.pyplot as plt
import mpld3
import streamlit.components.v1 as components
from helpers import *
import uuid
import os
import skimage.measure
st.set_page_config(layout="wide")
uploaded_file = st.file_uploader("Choose a file")
if uploaded_file is not None:
session_id = st.session_state.get('session_id', str(uuid.uuid4()))
session_id = st.text_area('session_id', value=session_id)
st.session_state['session_id'] = session_id
path = f'sessions/{session_id}'
if not(os.path.exists(path) and os.path.isdir(path)):
os.makedirs(path)
with open(f'sessions/{session_id}/image_orig.jpg', 'wb') as f:
f.write(uploaded_file.getvalue())
# Draw an image
img = cv2.imread(f'sessions/{session_id}/image_orig.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
while img.shape[0] > 1500 or img.shape[1] > 1500:
img = skimage.measure.block_reduce(img, (2,2,1), np.mean)
col1, col2, col3 = st.columns(3)
with col1:
fig = plt.figure()
plt.title("Original image")
img = np.uint8(img)
plt.imshow(img)
fig_html = mpld3.fig_to_html(fig)
components.html(fig_html,height=600)
img_for_save = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
cv2.imwrite(f"sessions/{session_id}/image.jpg", img_for_save)
# get course layer
get_separated_layers(session_id)
course = np.load(f'sessions/{session_id}/course_layer.npy')
with col2:
fig = plt.figure()
plt.title("Course layer")
plt.imshow(course, cmap='gray')
fig_html = mpld3.fig_to_html(fig)
components.html(fig_html,height=600)
get_runnable_mask(session_id)
runnable_mask = np.load(f'sessions/{session_id}/runnable_layer.npy')
with col3:
fig = plt.figure()
plt.title("Runnable area")
plt.imshow(runnable_mask * 255, cmap='gray')
fig_html = mpld3.fig_to_html(fig)
components.html(fig_html,height=600)
# detect triangle
col1, col2 = st.columns([1, 3])
with col1:
triangle_erode = st.slider(f'{session_id}: triangle_erode', 0, 20, 3)
triangle_dilate = st.slider(f'{session_id}: triangle_dilate', 0, 20, 5)
triangles = get_triangle(session_id, triangle_erode, triangle_dilate)
triangles = json.loads(st.text_area(f'{session_id}: triangle', json.dumps(triangles)))
with col2:
triangle_plot = np.zeros_like(course)
for triangle in triangles:
contours = np.array([[triangle['a'], triangle['b'], triangle['c']]], dtype=np.int32)
triangle_plot = cv2.drawContours(triangle_plot, contours, -1, 255, 3)
fig = plt.figure()
plt.title("Triangles detected")
plt.imshow(triangle_plot, cmap='gray')
fig_html = mpld3.fig_to_html(fig)
components.html(fig_html, height=600)
col1, col2 = st.columns([1, 3])
with col1:
circle_radius = st.slider(f'{session_id}: circle radius', 1, 59, 9, step=2)
circle_x = st.slider(f'{session_id}: circle X', circle_radius, img.shape[0]-circle_radius, circle_radius + (img.shape[0] - 2 * circle_radius) // 2)
circle_y = st.slider(f'{session_id}: circle Y', circle_radius, img.shape[1]-circle_radius, circle_radius + (img.shape[1] - 2 * circle_radius) // 2)
with col2:
fig = plt.figure()
plt.title("Original image")
img = np.uint8(img)
plt.imshow(img)
circle = plt.Circle((circle_x, circle_y), circle_radius, color='r')
plt.gca().add_patch(circle)
fig_html = mpld3.fig_to_html(fig)
components.html(fig_html, height=600)
# detect circles
col1, col2 = st.columns([1, 3])
generate_circles = None
with col1:
candidates_threshold = st.slider(f'{session_id}: candidates threshold', 0, 10 * circle_radius ** 2, int(0.8 * circle_radius ** 2))
match_threshold = st.slider(f'{session_id}: match threshold', 0, 10 * circle_radius ** 2, int(0.3 * circle_radius ** 2))
generate_circles = st.checkbox('generate circles (slow)')
if generate_circles:
circles = get_circles(session_id, circle_radius, candidates_threshold, match_threshold)
print(circles)
if f'{session_id}: circles' in st.session_state:
del st.session_state[f'{session_id}: circles']
else:
circles = {'circles': []}
circles = json.loads(st.text_area(f'{session_id}: circles', json.dumps(circles)))
with col2:
circles = circles['circles']
circles_plot = np.zeros_like(course)
for circle in circles:
circles_plot = cv2.circle(circles_plot, (circle['y'], circle['x']), circle['r'], 255, 3)
fig = plt.figure()
plt.title("Circles detected")
light_course = np.float32(course / 2)
plt.imshow(light_course + circles_plot, cmap='gray')
fig_html = mpld3.fig_to_html(fig)
components.html(fig_html, height=600)
# find configuration
points = []
for triangle in triangles:
points.append(np.int32((np.array(triangle['a']) + np.array(triangle['b']) + np.array(triangle['c'])) / 3)[0])
for circle in circles:
points.append(np.array([circle['x'], circle['y']]))
build_course_quadtree_graph(session_id)
build_course_distance_matrix(session_id, points)
if len(points) < 2 or len(circles) == 0:
st.warning('There are not enough points to run futher')
else:
col1, col2 = st.columns([1, 3])
with col1:
permutation = json.loads(st.text_area(f'{session_id}: permutation', json.dumps(get_circuit_permutation(session_id, points))))
with col2:
fig = plt.figure()
configuration = np.zeros((img.shape[0],img.shape[1],1), np.uint8)
for circle in circles:
cv2.circle(configuration, (circle['y'],circle['x']),12,255,10)
for triangle in triangles:
contours = np.array([[triangle['a'], triangle['b'], triangle['c']]], dtype=np.int32)
triangle_plot = cv2.drawContours(configuration, contours, -1, 255, 3)
for i in range(len(permutation) - 1):
cv2.line(configuration, (points[permutation[i]][1], points[permutation[i]][0]), ((points[permutation[(i+1)%len(permutation)]][1], points[permutation[(i+1)%len(permutation)]][0])), 255, 10)
plt.imshow(configuration + light_course.reshape((img.shape[0], img.shape[1], 1)), cmap='gray')
fig_html = mpld3.fig_to_html(fig)
components.html(fig_html, height=600)
build_runnable_quadtree_graph(session_id)
legs = []
for i in range(len(permutation)-1):
legs.append((points[permutation[i]], points[permutation[(i+1)]]))
get_routes(session_id, legs)
routes = np.load(f'sessions/{session_id}/routes_map.npy')
map_with_routes = (img / 255 - routes * 0.7)
fig = plt.figure()
plt.imshow(map_with_routes)
fig_html = mpld3.fig_to_html(fig)
components.html(fig_html, height=600)