Spaces:
Sleeping
Sleeping
Upload TheDistanceAssessor.py
Browse files- TheDistanceAssessor.py +921 -0
TheDistanceAssessor.py
ADDED
@@ -0,0 +1,921 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import cv2
|
2 |
+
import torch
|
3 |
+
import numpy as np
|
4 |
+
from skimage import morphology
|
5 |
+
import albumentations as A
|
6 |
+
import torch.nn.functional as F
|
7 |
+
import torch.nn as nn
|
8 |
+
from albumentations.pytorch import ToTensorV2
|
9 |
+
import matplotlib.pyplot as plt
|
10 |
+
from sklearn.linear_model import LinearRegression
|
11 |
+
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
|
12 |
+
import matplotlib.path as mplPath
|
13 |
+
import matplotlib.patches as patches
|
14 |
+
from ultralyticsplus import YOLO
|
15 |
+
|
16 |
+
def image_morpho(mask_prediction):
|
17 |
+
selem2 = morphology.disk(2)
|
18 |
+
closed = morphology.closing(mask_prediction, selem2)
|
19 |
+
|
20 |
+
return closed
|
21 |
+
|
22 |
+
def get_segformer_img(image_in, input_size=[224,224]):
|
23 |
+
transform_img = A.Compose([
|
24 |
+
A.Resize(height=input_size[0], width=input_size[1], interpolation=cv2.INTER_NEAREST),
|
25 |
+
A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),
|
26 |
+
ToTensorV2(p=1.0),
|
27 |
+
])
|
28 |
+
|
29 |
+
image_in = cv2.resize(image_in, (1920, 1080))
|
30 |
+
|
31 |
+
image_tr = transform_img(image=image_in)['image']
|
32 |
+
image_tr = image_tr.unsqueeze(0)
|
33 |
+
image_tr = image_tr.cpu()
|
34 |
+
|
35 |
+
return image_tr, image_in
|
36 |
+
|
37 |
+
def load_segformer(path_model):
|
38 |
+
|
39 |
+
model = torch.load(path_model, map_location=torch.device('cpu'))
|
40 |
+
model = model.cpu()
|
41 |
+
model.eval()
|
42 |
+
return model
|
43 |
+
|
44 |
+
def load_yolo(PATH_model):
|
45 |
+
model = YOLO(PATH_model)
|
46 |
+
|
47 |
+
model.overrides['conf'] = 0.25 # NMS confidence threshold
|
48 |
+
model.overrides['iou'] = 0.45 # NMS IoU threshold
|
49 |
+
model.overrides['agnostic_nms'] = False # NMS class-agnostic
|
50 |
+
model.overrides['max_det'] = 1000 # maximum number of detections per image
|
51 |
+
return model
|
52 |
+
|
53 |
+
def find_extreme_y_values(arr, values=[0, 6]):
|
54 |
+
"""
|
55 |
+
Optimized function to find the lowest and highest y-values (row indices) in a 2D array where 0 or 6 appears.
|
56 |
+
|
57 |
+
Parameters:
|
58 |
+
- arr: The input 2D NumPy array.
|
59 |
+
- values: The values to search for (default is [0, 6]).
|
60 |
+
|
61 |
+
Returns:
|
62 |
+
A tuple (lowest_y, highest_y) representing the lowest and highest y-values. If values are not found, returns None.
|
63 |
+
"""
|
64 |
+
mask = np.isin(arr, values)
|
65 |
+
rows_with_values = np.any(mask, axis=1)
|
66 |
+
|
67 |
+
y_indices = np.nonzero(rows_with_values)[0] # Directly finding non-zero (True) indices
|
68 |
+
|
69 |
+
if y_indices.size == 0:
|
70 |
+
return None, None # Early return if values not found
|
71 |
+
|
72 |
+
return y_indices[0], y_indices[-1]
|
73 |
+
|
74 |
+
def find_nearest_pairs(arr1, arr2):
|
75 |
+
# Convert lists to numpy arrays for vectorized operations
|
76 |
+
arr1_np = np.array(arr1)
|
77 |
+
arr2_np = np.array(arr2)
|
78 |
+
|
79 |
+
# Determine which array is shorter
|
80 |
+
if len(arr1_np) < len(arr2_np):
|
81 |
+
base_array, compare_array = arr1_np, arr2_np
|
82 |
+
else:
|
83 |
+
base_array, compare_array = arr2_np, arr1_np
|
84 |
+
|
85 |
+
paired_base = []
|
86 |
+
paired_compare = []
|
87 |
+
|
88 |
+
# Mask to keep track of paired elements
|
89 |
+
paired_mask = np.zeros(len(compare_array), dtype=bool)
|
90 |
+
|
91 |
+
for item in base_array:
|
92 |
+
# Calculate distances from the current item to all items in the compare_array
|
93 |
+
distances = np.linalg.norm(compare_array - item, axis=1)
|
94 |
+
nearest_index = np.argmin(distances)
|
95 |
+
paired_base.append(item)
|
96 |
+
paired_compare.append(compare_array[nearest_index])
|
97 |
+
# Mark the paired element to exclude it from further pairing
|
98 |
+
paired_mask[nearest_index] = True
|
99 |
+
|
100 |
+
# Check if all elements from the compare_array have been paired
|
101 |
+
if paired_mask.all():
|
102 |
+
break
|
103 |
+
|
104 |
+
paired_base = np.array(paired_base)
|
105 |
+
paired_compare = compare_array[paired_mask]
|
106 |
+
|
107 |
+
return (paired_base, paired_compare) if len(arr1_np) < len(arr2_np) else (paired_compare, paired_base)
|
108 |
+
|
109 |
+
def filter_crossings(image, edges_dict):
|
110 |
+
filtered_edges = {}
|
111 |
+
for key, values in edges_dict.items():
|
112 |
+
merged = [values[0]]
|
113 |
+
for start, end in values[1:]:
|
114 |
+
if start - merged[-1][1] < 50:
|
115 |
+
|
116 |
+
key_up = max([0, key-10])
|
117 |
+
key_down = min([image.shape[0]-1, key+10])
|
118 |
+
if key_up == 0:
|
119 |
+
key_up = key+20
|
120 |
+
if key_down == image.shape[0]-1:
|
121 |
+
key_down = key-20
|
122 |
+
|
123 |
+
edges_to_test_slope1 = robust_edges(image, [key_up], values=[0, 6], min_width=19)
|
124 |
+
edges_to_test_slope2 = robust_edges(image, [key_down], values=[0, 6], min_width=19)
|
125 |
+
|
126 |
+
values1, edges_to_test_slope1 = find_nearest_pairs(values, edges_to_test_slope1)
|
127 |
+
values2, edges_to_test_slope2 = find_nearest_pairs(values, edges_to_test_slope2)
|
128 |
+
|
129 |
+
differences_y = []
|
130 |
+
for i, value in enumerate(values1):
|
131 |
+
if start in value:
|
132 |
+
idx = list(value).index(start)
|
133 |
+
try:
|
134 |
+
differences_y.append(abs(start-edges_to_test_slope1[i][idx]))
|
135 |
+
except:
|
136 |
+
pass
|
137 |
+
if merged[-1][1] in value:
|
138 |
+
idx = list(value).index(merged[-1][1])
|
139 |
+
try:
|
140 |
+
differences_y.append(abs(merged[-1][1]-edges_to_test_slope1[i][idx]))
|
141 |
+
except:
|
142 |
+
pass
|
143 |
+
for i, value in enumerate(values2):
|
144 |
+
if start in value:
|
145 |
+
idx = list(value).index(start)
|
146 |
+
try:
|
147 |
+
differences_y.append(abs(start-edges_to_test_slope2[i][idx]))
|
148 |
+
except:
|
149 |
+
pass
|
150 |
+
if merged[-1][1] in value:
|
151 |
+
idx = list(value).index(merged[-1][1])
|
152 |
+
try:
|
153 |
+
differences_y.append(abs(merged[-1][1]-edges_to_test_slope2[i][idx]))
|
154 |
+
except:
|
155 |
+
pass
|
156 |
+
|
157 |
+
if any(element > 30 for element in differences_y):
|
158 |
+
merged[-1] = (merged[-1][0], end)
|
159 |
+
else:
|
160 |
+
merged.append((start, end))
|
161 |
+
else:
|
162 |
+
merged.append((start, end))
|
163 |
+
filtered_edges[key] = merged
|
164 |
+
|
165 |
+
return filtered_edges
|
166 |
+
|
167 |
+
def robust_edges(image, y_levels, values=[0, 6], min_width=19):
|
168 |
+
|
169 |
+
for y in y_levels:
|
170 |
+
row = image[y, :]
|
171 |
+
mask = np.isin(row, values).astype(int)
|
172 |
+
padded_mask = np.pad(mask, (1, 1), 'constant', constant_values=0)
|
173 |
+
diff = np.diff(padded_mask)
|
174 |
+
starts = np.where(diff == 1)[0]
|
175 |
+
ends = np.where(diff == -1)[0] - 1
|
176 |
+
|
177 |
+
# Filter sequences based on the minimum width criteria
|
178 |
+
filtered_edges = [(start, end) for start, end in zip(starts, ends) if end - start + 1 >= min_width]
|
179 |
+
filtered_edges = [(start, end) for start, end in filtered_edges if 0 not in (start, end) and 1919 not in (start, end)]
|
180 |
+
|
181 |
+
return filtered_edges
|
182 |
+
|
183 |
+
def find_edges(image, y_levels, values=[0, 6], min_width=19):
|
184 |
+
"""
|
185 |
+
Find start and end positions of continuous sequences of specified values at given y-levels in a 2D array,
|
186 |
+
filtering for sequences that meet or exceed a specified minimum width.
|
187 |
+
|
188 |
+
Parameters:
|
189 |
+
- arr: 2D NumPy array to search within.
|
190 |
+
- y_levels: List of y-levels (row indices) to examine.
|
191 |
+
- values: Values to search for (default is [0, 6]).
|
192 |
+
- min_width: Minimum width of sequences to be included in the results.
|
193 |
+
|
194 |
+
Returns:
|
195 |
+
A dict with y-levels as keys and lists of (start, end) tuples for each sequence found in that row that meets the width criteria.
|
196 |
+
"""
|
197 |
+
edges_dict = {}
|
198 |
+
for y in y_levels:
|
199 |
+
row = image[y, :]
|
200 |
+
mask = np.isin(row, values).astype(int)
|
201 |
+
padded_mask = np.pad(mask, (1, 1), 'constant', constant_values=0)
|
202 |
+
diff = np.diff(padded_mask)
|
203 |
+
starts = np.where(diff == 1)[0]
|
204 |
+
ends = np.where(diff == -1)[0] - 1
|
205 |
+
|
206 |
+
# Filter sequences based on the minimum width criteria
|
207 |
+
filtered_edges = [(start, end) for start, end in zip(starts, ends) if end - start + 1 >= min_width]
|
208 |
+
filtered_edges = [(start, end) for start, end in filtered_edges if 0 not in (start, end) and 1919 not in (start, end)]
|
209 |
+
|
210 |
+
edges_with_guard_rails = []
|
211 |
+
for edge in filtered_edges:
|
212 |
+
cutout_left = image[y,edge[0]-50:edge[0]][::-1]
|
213 |
+
cutout_right = image[y,edge[1]:edge[1]+50]
|
214 |
+
|
215 |
+
not_ones = np.where(cutout_left != 1)[0]
|
216 |
+
if len(not_ones) > 0 and not_ones[0] > 0:
|
217 |
+
last_one_index = not_ones[0] - 1
|
218 |
+
edge = (edge[0] - last_one_index,) + edge[1:]
|
219 |
+
else:
|
220 |
+
last_one_index = None if len(not_ones) == 0 else not_ones[-1] - 1
|
221 |
+
|
222 |
+
not_ones = np.where(cutout_right != 1)[0]
|
223 |
+
if len(not_ones) > 0 and not_ones[0] > 0:
|
224 |
+
last_one_index = not_ones[0] - 1
|
225 |
+
edge = (edge[0], edge[1] - last_one_index) + edge[2:]
|
226 |
+
else:
|
227 |
+
last_one_index = None if len(not_ones) == 0 else not_ones[-1] - 1
|
228 |
+
|
229 |
+
edges_with_guard_rails.append(edge)
|
230 |
+
|
231 |
+
edges_dict[y] = edges_with_guard_rails
|
232 |
+
|
233 |
+
edges_dict = {k: v for k, v in edges_dict.items() if v}
|
234 |
+
|
235 |
+
edges_dict = filter_crossings(image, edges_dict)
|
236 |
+
|
237 |
+
return edges_dict
|
238 |
+
|
239 |
+
def find_rails(arr, y_levels, values=[9, 10], min_width=5):
|
240 |
+
edges_all = []
|
241 |
+
for y in y_levels:
|
242 |
+
row = arr[y, :]
|
243 |
+
mask = np.isin(row, values).astype(int)
|
244 |
+
padded_mask = np.pad(mask, (1, 1), 'constant', constant_values=0)
|
245 |
+
diff = np.diff(padded_mask)
|
246 |
+
starts = np.where(diff == 1)[0]
|
247 |
+
ends = np.where(diff == -1)[0] - 1
|
248 |
+
|
249 |
+
# Filter sequences based on the minimum width criteria
|
250 |
+
filtered_edges = [(start, end) for start, end in zip(starts, ends) if end - start + 1 >= min_width]
|
251 |
+
filtered_edges = [(start, end) for start, end in filtered_edges if 0 not in (start, end) and 1919 not in (start, end)]
|
252 |
+
edges_all = filtered_edges
|
253 |
+
|
254 |
+
return edges_all
|
255 |
+
|
256 |
+
def mark_edges(arr, edges_dict, mark_value=30):
|
257 |
+
"""
|
258 |
+
Marks a 5x5 zone around the edges found in the array with a specific value.
|
259 |
+
|
260 |
+
Parameters:
|
261 |
+
- arr: The original 2D NumPy array.
|
262 |
+
- edges_dict: A dictionary with y-levels as keys and lists of (start, end) tuples for edges.
|
263 |
+
- mark_value: The value used to mark the edges.
|
264 |
+
|
265 |
+
Returns:
|
266 |
+
The modified array with marked zones.
|
267 |
+
"""
|
268 |
+
marked_arr = np.copy(arr) # Create a copy of the array to avoid modifying the original
|
269 |
+
offset = 2 # To mark a 5x5 area, we go 2 pixels in each direction from the center
|
270 |
+
|
271 |
+
for y, edges in edges_dict.items():
|
272 |
+
for start, end in edges:
|
273 |
+
# Mark a 5x5 zone around the start and end positions
|
274 |
+
for dy in range(-offset, offset + 1):
|
275 |
+
for dx in range(-offset, offset + 1):
|
276 |
+
# Check array bounds before marking
|
277 |
+
if 0 <= y + dy < marked_arr.shape[0] and 0 <= start + dx < marked_arr.shape[1]:
|
278 |
+
marked_arr[y + dy, start + dx] = mark_value
|
279 |
+
if 0 <= y + dy < marked_arr.shape[0] and 0 <= end + dx < marked_arr.shape[1]:
|
280 |
+
marked_arr[y + dy, end + dx] = mark_value
|
281 |
+
|
282 |
+
return marked_arr
|
283 |
+
|
284 |
+
def find_rail_sides(img, edges_dict):
|
285 |
+
left_border = []
|
286 |
+
right_border = []
|
287 |
+
for y,xs in edges_dict.items():
|
288 |
+
rails = find_rails(img, [y], values=[9,10], min_width=5)
|
289 |
+
left_border_actual = [min(xs)[0],y]
|
290 |
+
right_border_actual = [max(xs)[1],y]
|
291 |
+
|
292 |
+
for zone in rails:
|
293 |
+
if abs(zone[1]-left_border_actual[0]) < y*0.04: # dynamic treshold
|
294 |
+
left_border_actual[0] = zone[0]
|
295 |
+
if abs(zone[0]-right_border_actual[0]) < y*0.04:
|
296 |
+
right_border_actual[0] = zone[1]
|
297 |
+
|
298 |
+
left_border.append(left_border_actual)
|
299 |
+
right_border.append(right_border_actual)
|
300 |
+
|
301 |
+
# removing detected uncontioussness
|
302 |
+
left_border, flags_l, _ = robust_rail_sides(left_border) # filter outliers
|
303 |
+
right_border, flags_r, _ = robust_rail_sides(right_border)
|
304 |
+
|
305 |
+
return left_border, right_border, flags_l, flags_r
|
306 |
+
|
307 |
+
def robust_rail_sides(border, threshold=7):
|
308 |
+
border = np.array(border)
|
309 |
+
if border.size > 0:
|
310 |
+
# delete borders found on the bottom side of the image
|
311 |
+
border = border[border[:, 1] != 1079]
|
312 |
+
|
313 |
+
steps_x = np.diff(border[:, 0])
|
314 |
+
median_step = np.median(np.abs(steps_x))
|
315 |
+
|
316 |
+
threshold_step = np.abs(threshold*np.abs(median_step))
|
317 |
+
treshold_overcommings = abs(steps_x) > abs(threshold_step)
|
318 |
+
|
319 |
+
flags = []
|
320 |
+
|
321 |
+
if True not in treshold_overcommings:
|
322 |
+
return border, flags, []
|
323 |
+
else:
|
324 |
+
overcommings_indices = [i for i, element in enumerate(treshold_overcommings) if element == True]
|
325 |
+
if overcommings_indices and np.all(np.diff(overcommings_indices) == 1):
|
326 |
+
overcommings_indices = [overcommings_indices[0]]
|
327 |
+
|
328 |
+
filtered_border = border
|
329 |
+
|
330 |
+
previously_deleted = []
|
331 |
+
for i in overcommings_indices:
|
332 |
+
for item in previously_deleted:
|
333 |
+
if item[0] < i:
|
334 |
+
i -= item[1]
|
335 |
+
first_part = filtered_border[:i+1]
|
336 |
+
second_part = filtered_border[i+1:]
|
337 |
+
if len(second_part)<2:
|
338 |
+
filtered_border = first_part
|
339 |
+
previously_deleted.append([i,len(second_part)])
|
340 |
+
elif len(first_part)<2:
|
341 |
+
filtered_border = second_part
|
342 |
+
previously_deleted.append([i,len(first_part)])
|
343 |
+
else:
|
344 |
+
first_b, _, deleted_first = robust_rail_sides(first_part)
|
345 |
+
second_b, _, _ = robust_rail_sides(second_part)
|
346 |
+
filtered_border = np.concatenate((first_b,second_b), axis=0)
|
347 |
+
|
348 |
+
if deleted_first:
|
349 |
+
for deleted_item in deleted_first:
|
350 |
+
if deleted_item[0]<=i:
|
351 |
+
i -= deleted_item[1]
|
352 |
+
|
353 |
+
flags.append(i)
|
354 |
+
return filtered_border, flags, previously_deleted
|
355 |
+
else:
|
356 |
+
return border, [], []
|
357 |
+
|
358 |
+
def find_dist_from_edges(id_map, image, edges_dict, left_border, right_border, real_life_width_mm, real_life_target_mm, mark_value=30):
|
359 |
+
"""
|
360 |
+
Mark regions representing a real-life distance (e.g., 2 meters) to the left and right from the furthest edges.
|
361 |
+
|
362 |
+
Parameters:
|
363 |
+
- arr: 2D NumPy array representing the id_map.
|
364 |
+
- edges_dict: Dictionary with y-levels as keys and lists of (start, end) tuples for edges.
|
365 |
+
- real_life_width_mm: The real-world width in millimeters that the average sequence width represents.
|
366 |
+
- real_life_target_mm: The real-world distance in millimeters to mark from the edges.
|
367 |
+
|
368 |
+
Returns:
|
369 |
+
- A NumPy array with the marked regions.
|
370 |
+
"""
|
371 |
+
# Calculate the rail widths
|
372 |
+
diffs_widths = {k: sum(e-s for s, e in v) / len(v) for k, v in edges_dict.items() if v}
|
373 |
+
diffs_width = {k: max(e-s for s, e in v) for k, v in edges_dict.items() if v}
|
374 |
+
|
375 |
+
# Pixel to mm scale factor
|
376 |
+
scale_factors = {k: real_life_width_mm / v for k, v in diffs_width.items()}
|
377 |
+
# Converting the real-life target distance to pixels
|
378 |
+
target_distances_px = {k: int(real_life_target_mm / v) for k, v in scale_factors.items()}
|
379 |
+
|
380 |
+
# Mark the regions representing the target distance to the left and right from the furthest edges
|
381 |
+
end_points_left = {}
|
382 |
+
region_levels_left = []
|
383 |
+
for point in left_border:
|
384 |
+
min_edge = point[0]
|
385 |
+
|
386 |
+
# Ensure we stay within the image bounds
|
387 |
+
#left_mark_start = max(0, min_edge - int(target_distances_px[point[1]]))
|
388 |
+
left_mark_start = min_edge - int(target_distances_px[point[1]])
|
389 |
+
end_points_left[point[1]] = left_mark_start
|
390 |
+
|
391 |
+
# Left region points
|
392 |
+
if left_mark_start < min_edge:
|
393 |
+
y_values = np.arange(left_mark_start, min_edge)
|
394 |
+
x_values = np.full_like(y_values, point[1])
|
395 |
+
region_line = np.column_stack((x_values, y_values))
|
396 |
+
region_levels_left.append(region_line)
|
397 |
+
|
398 |
+
end_points_right = {}
|
399 |
+
region_levels_right = []
|
400 |
+
for point in right_border:
|
401 |
+
max_edge = point[0]
|
402 |
+
|
403 |
+
# Ensure we stay within the image bounds
|
404 |
+
right_mark_end = min(id_map.shape[1], max_edge + int(target_distances_px[point[1]]))
|
405 |
+
if right_mark_end != id_map.shape[1]:
|
406 |
+
end_points_right[point[1]] = right_mark_end
|
407 |
+
|
408 |
+
# Right region points
|
409 |
+
if max_edge < right_mark_end:
|
410 |
+
y_values = np.arange(max_edge, right_mark_end)
|
411 |
+
x_values = np.full_like(y_values, point[1])
|
412 |
+
region_line = np.column_stack((x_values, y_values))
|
413 |
+
region_levels_right.append(region_line)
|
414 |
+
|
415 |
+
return id_map, end_points_left, end_points_right, region_levels_left, region_levels_right
|
416 |
+
|
417 |
+
def bresenham_line(x0, y0, x1, y1):
|
418 |
+
"""
|
419 |
+
Generate the coordinates of a line from (x0, y0) to (x1, y1) using Bresenham's algorithm.
|
420 |
+
"""
|
421 |
+
line = []
|
422 |
+
dx = abs(x1 - x0)
|
423 |
+
dy = -abs(y1 - y0)
|
424 |
+
sx = 1 if x0 < x1 else -1
|
425 |
+
sy = 1 if y0 < y1 else -1
|
426 |
+
err = dx + dy # error value e_xy
|
427 |
+
|
428 |
+
while True:
|
429 |
+
line.append((x0, y0)) # Add the current point to the line
|
430 |
+
if x0 == x1 and y0 == y1:
|
431 |
+
break
|
432 |
+
e2 = 2 * err
|
433 |
+
if e2 >= dy: # e_xy+e_x > 0
|
434 |
+
err += dy
|
435 |
+
x0 += sx
|
436 |
+
if e2 <= dx: # e_xy+e_y < 0
|
437 |
+
err += dx
|
438 |
+
y0 += sy
|
439 |
+
|
440 |
+
return line
|
441 |
+
|
442 |
+
def interpolate_end_points(end_points_dict, flags):
|
443 |
+
line_arr = []
|
444 |
+
ys = list(end_points_dict.keys())
|
445 |
+
xs = list(end_points_dict.values())
|
446 |
+
|
447 |
+
if flags and len(flags) == 1:
|
448 |
+
pass
|
449 |
+
elif flags and np.all(np.diff(flags) == 1):
|
450 |
+
flags = [flags[0]]
|
451 |
+
|
452 |
+
for i in range(0, len(ys) - 1):
|
453 |
+
if i in flags:
|
454 |
+
continue
|
455 |
+
y1, y2 = ys[i], ys[i + 1]
|
456 |
+
x1, x2 = xs[i], xs[i + 1]
|
457 |
+
line = np.array(bresenham_line(x1, y1, x2, y2))
|
458 |
+
if np.any(line[:, 0] < 0):
|
459 |
+
line = line[line[:, 0] > 0]
|
460 |
+
line_arr = line_arr + list(line)
|
461 |
+
|
462 |
+
return line_arr
|
463 |
+
|
464 |
+
def extrapolate_line(pixels, image, min_y=None, extr_pixels=10):
|
465 |
+
"""
|
466 |
+
Extrapolate a line based on the last segment using linear regression.
|
467 |
+
|
468 |
+
Parameters:
|
469 |
+
- pixels: List of (x, y) tuples representing line pixel coordinates.
|
470 |
+
- image: 2D numpy array representing the image.
|
471 |
+
- min_y: Minimum y-value to extrapolate to (optional).
|
472 |
+
|
473 |
+
Returns:
|
474 |
+
- A list of new extrapolated (x, y) pixel coordinates.
|
475 |
+
"""
|
476 |
+
if len(pixels) < extr_pixels:
|
477 |
+
print("Not enough pixels to perform extrapolation.")
|
478 |
+
return []
|
479 |
+
|
480 |
+
recent_pixels = np.array(pixels[-extr_pixels:])
|
481 |
+
|
482 |
+
X = recent_pixels[:, 0].reshape(-1, 1) # Reshape for sklearn
|
483 |
+
y = recent_pixels[:, 1]
|
484 |
+
|
485 |
+
model = LinearRegression()
|
486 |
+
model.fit(X, y)
|
487 |
+
|
488 |
+
slope = model.coef_[0]
|
489 |
+
intercept = model.intercept_
|
490 |
+
|
491 |
+
extrapolate = lambda x: slope * x + intercept
|
492 |
+
|
493 |
+
# Calculate direction based on last two pixels
|
494 |
+
dx, dy = 0, 0 # Default values
|
495 |
+
|
496 |
+
x_diffs = []
|
497 |
+
y_diffs = []
|
498 |
+
for i in range(1,extr_pixels-1):
|
499 |
+
x_diffs.append(pixels[-i][0] - pixels[-(i+1)][0])
|
500 |
+
y_diffs.append(pixels[-i][1] - pixels[-(i+1)][1])
|
501 |
+
|
502 |
+
x_diff = x_diffs[np.argmax(np.abs(x_diffs))]
|
503 |
+
y_diff = y_diffs[np.argmax(np.abs(y_diffs))]
|
504 |
+
|
505 |
+
if abs(int(x_diff)) >= abs(int(y_diff)):
|
506 |
+
dx = 1 if x_diff >= 0 else -1
|
507 |
+
else:
|
508 |
+
dy = 1 if y_diff >= 0 else -1
|
509 |
+
|
510 |
+
last_pixel = pixels[-1]
|
511 |
+
new_pixels = []
|
512 |
+
x, y = last_pixel
|
513 |
+
|
514 |
+
min_y = min_y if min_y is not None else image.shape[0] - 1
|
515 |
+
|
516 |
+
while 0 <= x < image.shape[1] and min_y <= y < image.shape[0]:
|
517 |
+
if dx != 0: # Horizontal or diagonal movement
|
518 |
+
x += dx
|
519 |
+
y = int(extrapolate(x))
|
520 |
+
elif dy != 0: # Vertical movement
|
521 |
+
y += dy
|
522 |
+
# For vertical lines, approximate x based on the last known value
|
523 |
+
x = int(x)
|
524 |
+
|
525 |
+
if 0 <= y < image.shape[0] and 0 <= x < image.shape[1]:
|
526 |
+
new_pixels.append((x, y))
|
527 |
+
else:
|
528 |
+
break
|
529 |
+
|
530 |
+
return new_pixels
|
531 |
+
|
532 |
+
def extrapolate_borders(dist_marked_id_map, border_l, border_r, lowest_y):
|
533 |
+
|
534 |
+
#border_extrapolation_l1 = extrapolate_line(border_l, dist_marked_id_map, lowest_y)
|
535 |
+
border_extrapolation_l2 = extrapolate_line(border_l[::-1], dist_marked_id_map, lowest_y)
|
536 |
+
|
537 |
+
#border_extrapolation_r1 = extrapolate_line(border_r, dist_marked_id_map, lowest_y)
|
538 |
+
border_extrapolation_r2 = extrapolate_line(border_r[::-1], dist_marked_id_map, lowest_y)
|
539 |
+
|
540 |
+
#border_l = border_extrapolation_l2[::-1] + border_l + border_extrapolation_l1
|
541 |
+
#border_r = border_extrapolation_r2[::-1] + border_r + border_extrapolation_r1
|
542 |
+
|
543 |
+
border_l = border_extrapolation_l2[::-1] + border_l
|
544 |
+
border_r = border_extrapolation_r2[::-1] + border_r
|
545 |
+
|
546 |
+
return border_l, border_r
|
547 |
+
|
548 |
+
def find_zone_border(id_map, image, edges, irl_width_mm=1435, irl_target_mm=1000, lowest_y = 0):
|
549 |
+
|
550 |
+
left_border, right_border, flags_l, flags_r = find_rail_sides(id_map, edges)
|
551 |
+
|
552 |
+
dist_marked_id_map, end_points_left, end_points_right, left_region, right_region = find_dist_from_edges(id_map, image, edges, left_border, right_border, irl_width_mm, irl_target_mm)
|
553 |
+
|
554 |
+
border_l = interpolate_end_points(end_points_left, flags_l)
|
555 |
+
border_r = interpolate_end_points(end_points_right, flags_r)
|
556 |
+
|
557 |
+
border_l, border_r = extrapolate_borders(dist_marked_id_map, border_l, border_r, lowest_y)
|
558 |
+
|
559 |
+
return [border_l, border_r],[left_region, right_region]
|
560 |
+
|
561 |
+
def get_clues(segmentation_mask, number_of_clues):
|
562 |
+
|
563 |
+
lowest, highest = find_extreme_y_values(segmentation_mask)
|
564 |
+
if lowest is not None and highest is not None:
|
565 |
+
clue_step = int((highest - lowest) / number_of_clues+1)
|
566 |
+
clues = []
|
567 |
+
for i in range(number_of_clues):
|
568 |
+
clues.append(highest - (i*clue_step))
|
569 |
+
clues.append(lowest+int(0.5*clue_step))
|
570 |
+
|
571 |
+
return clues
|
572 |
+
else:
|
573 |
+
return []
|
574 |
+
|
575 |
+
def border_handler(id_map, image, edges, target_distances):
|
576 |
+
|
577 |
+
lowest, _ = find_extreme_y_values(id_map)
|
578 |
+
borders = []
|
579 |
+
regions = []
|
580 |
+
for target in target_distances:
|
581 |
+
borders_regions = find_zone_border(id_map, image, edges, irl_target_mm=target, lowest_y = lowest)
|
582 |
+
borders.append(borders_regions[0])
|
583 |
+
regions.append(borders_regions[1])
|
584 |
+
|
585 |
+
return borders, id_map, regions
|
586 |
+
|
587 |
+
def segment(input_image, model_seg, image_size):
|
588 |
+
image_norm, image = get_segformer_img(input_image, image_size)
|
589 |
+
|
590 |
+
outputs = model_seg(image_norm)
|
591 |
+
|
592 |
+
logits = outputs.logits
|
593 |
+
upsampled_logits = nn.functional.interpolate(
|
594 |
+
logits,
|
595 |
+
size=image_norm.shape[-2:],
|
596 |
+
mode="bilinear",
|
597 |
+
align_corners=False
|
598 |
+
)
|
599 |
+
|
600 |
+
output = upsampled_logits.float()
|
601 |
+
|
602 |
+
confidence_scores = F.softmax(output, dim=1).cpu().detach().numpy().squeeze()
|
603 |
+
id_map = np.argmax(confidence_scores, axis=0).astype(np.uint8)
|
604 |
+
id_map = image_morpho(id_map)
|
605 |
+
|
606 |
+
id_map = cv2.resize(id_map, [input_image.shape[1],input_image.shape[0]], interpolation=cv2.INTER_NEAREST)
|
607 |
+
return id_map, image
|
608 |
+
|
609 |
+
def detect(model_det, image):
|
610 |
+
|
611 |
+
results = model_det.predict(image)
|
612 |
+
|
613 |
+
return results, model_det, image
|
614 |
+
|
615 |
+
def manage_detections(results, model):
|
616 |
+
bbox = results[0].boxes.xywh.tolist()
|
617 |
+
cls = results[0].boxes.cls.tolist()
|
618 |
+
accepted_stationary = np.array([24,25,28,36])
|
619 |
+
accepted_moving = np.array([0,1,2,3,7,15,16,17,18,19])
|
620 |
+
boxes_moving = {}
|
621 |
+
boxes_stationary = {}
|
622 |
+
if len(bbox) > 0:
|
623 |
+
for xywh, clss in zip(bbox, cls):
|
624 |
+
if clss in accepted_moving:
|
625 |
+
if clss in boxes_moving.keys() and len(boxes_moving[clss]) > 0:
|
626 |
+
boxes_moving[clss].append(xywh)
|
627 |
+
else:
|
628 |
+
boxes_moving[clss] = [xywh]
|
629 |
+
if clss in accepted_stationary:
|
630 |
+
if clss in boxes_stationary.keys() and len(boxes_stationary[clss]) > 0:
|
631 |
+
boxes_stationary[clss].append(xywh)
|
632 |
+
else:
|
633 |
+
boxes_stationary[clss] = [xywh]
|
634 |
+
|
635 |
+
return boxes_moving, boxes_stationary
|
636 |
+
|
637 |
+
def compute_detection_borders(borders, output_dims=[1080,1920]):
|
638 |
+
det_height = output_dims[0]-1
|
639 |
+
det_width = output_dims[1]-1
|
640 |
+
|
641 |
+
for i,border in enumerate(borders):
|
642 |
+
border_l = np.array(border[0])
|
643 |
+
|
644 |
+
if list(border_l):
|
645 |
+
pass
|
646 |
+
else:
|
647 |
+
border_l=np.array([[0,0],[0,0]])
|
648 |
+
|
649 |
+
endpoints_l = [border_l[0],border_l[-1]]
|
650 |
+
|
651 |
+
border_r = np.array(border[1])
|
652 |
+
if list(border_r):
|
653 |
+
pass
|
654 |
+
else:
|
655 |
+
border_r=np.array([[0,0],[0,0]])
|
656 |
+
|
657 |
+
endpoints_r = [border_r[0],border_r[-1]]
|
658 |
+
|
659 |
+
if np.array_equal(np.array([[0,0],[0,0]]), endpoints_l):
|
660 |
+
endpoints_l = [[0,endpoints_r[0][1]],[0,endpoints_r[1][1]]]
|
661 |
+
|
662 |
+
if np.array_equal(np.array([[0,0],[0,0]]), endpoints_r):
|
663 |
+
endpoints_r = [[det_width,endpoints_l[0][1]],[det_width,endpoints_l[1][1]]]
|
664 |
+
|
665 |
+
interpolated_top = bresenham_line(endpoints_l[1][0],endpoints_l[1][1],endpoints_r[1][0],endpoints_r[1][1])
|
666 |
+
|
667 |
+
zero_range = [0,1,2,3]
|
668 |
+
height_range = [det_height,det_height-1,det_height-2,det_height-3]
|
669 |
+
width_range = [det_width,det_width-1,det_width-2,det_width-3]
|
670 |
+
|
671 |
+
if (endpoints_l[0][0] in zero_range and endpoints_r[0][1] in height_range):
|
672 |
+
y_values = np.arange(endpoints_l[0][1], det_height)
|
673 |
+
x_values = np.full_like(y_values, 0)
|
674 |
+
bottom1 = np.column_stack((x_values, y_values))
|
675 |
+
|
676 |
+
x_values = np.arange(0, endpoints_r[0][0])
|
677 |
+
y_values = np.full_like(x_values, det_height)
|
678 |
+
bottom2 = np.column_stack((x_values, y_values))
|
679 |
+
|
680 |
+
interpolated_bottom = np.vstack((bottom1, bottom2))
|
681 |
+
|
682 |
+
elif (endpoints_l[0][1] in height_range and endpoints_r[0][0] in width_range):
|
683 |
+
y_values = np.arange(endpoints_r[0][1], det_height)
|
684 |
+
x_values = np.full_like(y_values, det_width)
|
685 |
+
bottom1 = np.column_stack((x_values, y_values))
|
686 |
+
|
687 |
+
x_values = np.arange(endpoints_l[0][0], det_width)
|
688 |
+
y_values = np.full_like(x_values, det_height)
|
689 |
+
bottom2 = np.column_stack((x_values, y_values))
|
690 |
+
|
691 |
+
interpolated_bottom = np.vstack((bottom1, bottom2))
|
692 |
+
|
693 |
+
elif endpoints_l[0][0] in zero_range and endpoints_r[0][0] in width_range:
|
694 |
+
y_values = np.arange(endpoints_l[0][1], det_height)
|
695 |
+
x_values = np.full_like(y_values, 0)
|
696 |
+
bottom1 = np.column_stack((x_values, y_values))
|
697 |
+
|
698 |
+
y_values = np.arange(endpoints_r[0][1], det_height)
|
699 |
+
x_values = np.full_like(y_values, det_width)
|
700 |
+
bottom2 = np.column_stack((x_values, y_values))
|
701 |
+
|
702 |
+
bottom3_mid = bresenham_line(bottom1[-1][0],bottom1[-1][1],bottom2[-1][0],bottom2[-1][1])
|
703 |
+
|
704 |
+
interpolated_bottom = np.vstack((bottom1, bottom2, bottom3_mid))
|
705 |
+
|
706 |
+
|
707 |
+
else:
|
708 |
+
interpolated_bottom = bresenham_line(endpoints_l[0][0],endpoints_l[0][1],endpoints_r[0][0],endpoints_r[0][1])
|
709 |
+
|
710 |
+
borders[i].append(interpolated_bottom)
|
711 |
+
borders[i].append(interpolated_top)
|
712 |
+
|
713 |
+
return borders
|
714 |
+
|
715 |
+
def get_bounding_box_points(cx, cy, w, h):
|
716 |
+
top_left = (cx - w / 2, cy - h / 2)
|
717 |
+
top_right = (cx + w / 2, cy - h / 2)
|
718 |
+
bottom_right = (cx + w / 2, cy + h / 2)
|
719 |
+
bottom_left = (cx - w / 2, cy + h / 2)
|
720 |
+
|
721 |
+
corners = [top_left, top_right, bottom_right, bottom_left]
|
722 |
+
|
723 |
+
def interpolate(point1, point2, fraction):
|
724 |
+
"""Interpolate between two points at a given fraction of the distance."""
|
725 |
+
return (point1[0] + fraction * (point2[0] - point1[0]),
|
726 |
+
point1[1] + fraction * (point2[1] - point1[1]))
|
727 |
+
|
728 |
+
points = []
|
729 |
+
for i in range(4):
|
730 |
+
next_i = (i + 1) % 4
|
731 |
+
points.append(corners[i])
|
732 |
+
points.append(interpolate(corners[i], corners[next_i], 1 / 3))
|
733 |
+
points.append(interpolate(corners[i], corners[next_i], 2 / 3))
|
734 |
+
|
735 |
+
return points
|
736 |
+
|
737 |
+
def classify_detections(boxes_moving, boxes_stationary, borders, img_dims, output_dims=[1080,1920]):
|
738 |
+
img_h, img_w, _ = img_dims
|
739 |
+
img_h_scaletofullHD = output_dims[1]/img_w
|
740 |
+
img_w_scaletofullHD = output_dims[0]/img_h
|
741 |
+
colors = ["yellow","orange","red","green","blue"]
|
742 |
+
|
743 |
+
borders = compute_detection_borders(borders,output_dims)
|
744 |
+
|
745 |
+
boxes_info = []
|
746 |
+
|
747 |
+
if boxes_moving or boxes_stationary:
|
748 |
+
if boxes_moving:
|
749 |
+
for item, coords in boxes_moving.items():
|
750 |
+
for coord in coords:
|
751 |
+
x = coord[0]*img_w_scaletofullHD
|
752 |
+
y = coord[1]*img_h_scaletofullHD
|
753 |
+
w = coord[2]*img_w_scaletofullHD
|
754 |
+
h = coord[3]*img_h_scaletofullHD
|
755 |
+
|
756 |
+
points_to_test = get_bounding_box_points(x, y, w, h)
|
757 |
+
|
758 |
+
complete_border = []
|
759 |
+
criticality = -1
|
760 |
+
color = None
|
761 |
+
for i,border in enumerate(reversed(borders)):
|
762 |
+
border_nonempty = [np.array(arr) for arr in border if np.array(arr).size > 0]
|
763 |
+
complete_border = np.vstack((border_nonempty))
|
764 |
+
instance_border_path = mplPath.Path(np.array(complete_border))
|
765 |
+
|
766 |
+
is_inside_borders = False
|
767 |
+
for point in points_to_test:
|
768 |
+
is_inside = instance_border_path.contains_point(point)
|
769 |
+
if is_inside:
|
770 |
+
is_inside_borders = True
|
771 |
+
|
772 |
+
if is_inside_borders:
|
773 |
+
criticality = i
|
774 |
+
color = colors[i]
|
775 |
+
|
776 |
+
if criticality == -1:
|
777 |
+
color = colors[3]
|
778 |
+
|
779 |
+
boxes_info.append([item, criticality, color, [x, y], [w, h], 1])
|
780 |
+
|
781 |
+
if boxes_stationary:
|
782 |
+
for item, coords in boxes_stationary.items():
|
783 |
+
for coord in coords:
|
784 |
+
x = coord[0]*img_w_scaletofullHD
|
785 |
+
y = coord[1]*img_h_scaletofullHD
|
786 |
+
w = coord[2]*img_w_scaletofullHD
|
787 |
+
h = coord[3]*img_h_scaletofullHD
|
788 |
+
|
789 |
+
points_to_test = get_bounding_box_points(x, y, w, h)
|
790 |
+
|
791 |
+
complete_border = []
|
792 |
+
criticality = -1
|
793 |
+
color = None
|
794 |
+
is_inside_borders = 0
|
795 |
+
for i,border in enumerate(reversed(borders), start=len(borders) - 1):
|
796 |
+
border_nonempty = [np.array(arr) for arr in border if np.array(arr).size > 0]
|
797 |
+
complete_border = np.vstack(border_nonempty)
|
798 |
+
instance_border_path = mplPath.Path(np.array(complete_border))
|
799 |
+
|
800 |
+
is_inside_borders = False
|
801 |
+
for point in points_to_test:
|
802 |
+
is_inside = instance_border_path.contains_point(point)
|
803 |
+
if is_inside:
|
804 |
+
is_inside_borders = True
|
805 |
+
|
806 |
+
if is_inside_borders:
|
807 |
+
criticality = i
|
808 |
+
color = colors[4]
|
809 |
+
|
810 |
+
if criticality == -1:
|
811 |
+
color = colors[3]
|
812 |
+
|
813 |
+
boxes_info.append([item, criticality, color, [x, y], [w, h], 0])
|
814 |
+
|
815 |
+
return boxes_info
|
816 |
+
|
817 |
+
else:
|
818 |
+
print("No accepted detections in this image.")
|
819 |
+
return []
|
820 |
+
|
821 |
+
def draw_classification(classification, id_map):
|
822 |
+
if classification:
|
823 |
+
for box in classification:
|
824 |
+
x,y = box[3]
|
825 |
+
mark_value = 30
|
826 |
+
|
827 |
+
x_start = int(max(x - 2, 0))
|
828 |
+
x_end = int(min(x + 3, id_map.shape[1]))
|
829 |
+
y_start = int(max(y - 2, 0))
|
830 |
+
y_end = int(min(y + 3, id_map.shape[0]))
|
831 |
+
|
832 |
+
id_map[y_start:y_end, x_start:x_end] = mark_value
|
833 |
+
else:
|
834 |
+
return
|
835 |
+
|
836 |
+
def get_result(classification, id_map, names, borders, image, regions):
|
837 |
+
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
838 |
+
image = cv2.resize(image, (id_map.shape[1], id_map.shape[0]), interpolation = cv2.INTER_LINEAR)
|
839 |
+
fig = plt.figure(figsize=(16, 9), dpi=100)
|
840 |
+
plt.imshow(image, cmap='gray')
|
841 |
+
|
842 |
+
if classification:
|
843 |
+
for box in classification:
|
844 |
+
|
845 |
+
boxes = True
|
846 |
+
cx,cy = box[3]
|
847 |
+
name = names[box[0]]
|
848 |
+
if boxes:
|
849 |
+
w,h = box[4]
|
850 |
+
x = cx - w / 2
|
851 |
+
y = cy - h / 2
|
852 |
+
rect = patches.Rectangle((x, y), w, h, linewidth=2, edgecolor=box[2], facecolor='none')
|
853 |
+
|
854 |
+
ax = plt.gca()
|
855 |
+
ax.add_patch(rect)
|
856 |
+
plt.text(x, y-17, name, color='black', fontsize=10, ha='center', va='center', fontweight='bold', bbox=dict(facecolor=box[2], edgecolor='none', alpha=1))
|
857 |
+
else:
|
858 |
+
plt.imshow(id_map, cmap='gray')
|
859 |
+
plt.text(cx, cy+10, name, color=box[2], fontsize=10, ha='center', va='center', fontweight='bold')
|
860 |
+
|
861 |
+
for region in regions:
|
862 |
+
for side in region:
|
863 |
+
for line in side:
|
864 |
+
line = np.array(line)
|
865 |
+
plt.plot(line[:,1], line[:,0] ,'-', color='lightgrey', marker=None, linewidth=0.5)
|
866 |
+
plt.ylim(0, 1080)
|
867 |
+
plt.xlim(0, 1920)
|
868 |
+
plt.gca().invert_yaxis()
|
869 |
+
|
870 |
+
colors = ['yellow','orange','red']
|
871 |
+
borders.reverse()
|
872 |
+
for i,border in enumerate(borders):
|
873 |
+
for side in border:
|
874 |
+
side = np.array(side)
|
875 |
+
if side.size > 0:
|
876 |
+
plt.plot(side[:,0],side[:,1] ,'-', color=colors[i], marker=None, linewidth=0.6) #color=colors[i]
|
877 |
+
plt.ylim(0, 1080)
|
878 |
+
plt.xlim(0, 1920)
|
879 |
+
plt.gca().invert_yaxis()
|
880 |
+
|
881 |
+
plt.tight_layout()
|
882 |
+
canvas = FigureCanvas(fig)
|
883 |
+
canvas.draw()
|
884 |
+
width, height = fig.get_size_inches() * fig.get_dpi()
|
885 |
+
image = np.frombuffer(canvas.tostring_rgb(), dtype='uint8').reshape(int(height), int(width), 3)
|
886 |
+
|
887 |
+
plt.close(fig) # Close the figure to free memory
|
888 |
+
|
889 |
+
return image
|
890 |
+
|
891 |
+
def run(input_image, model_seg, model_det, image_size, target_distances, num_ys = 10):
|
892 |
+
|
893 |
+
segmentation_mask, image = segment(input_image, model_seg, image_size)
|
894 |
+
|
895 |
+
# Border search
|
896 |
+
clues = get_clues(segmentation_mask, num_ys)
|
897 |
+
edges = find_edges(segmentation_mask, clues, min_width=0)
|
898 |
+
borders, id_map, regions = border_handler(segmentation_mask, image, edges, target_distances)
|
899 |
+
|
900 |
+
# Detection
|
901 |
+
results, model, image = detect(model_det, input_image)
|
902 |
+
boxes_moving, boxes_stationary = manage_detections(results, model)
|
903 |
+
|
904 |
+
classification = classify_detections(boxes_moving, boxes_stationary, borders, image.shape, output_dims=segmentation_mask.shape)
|
905 |
+
|
906 |
+
output_image = get_result(classification, id_map, model.names, borders, image, regions)
|
907 |
+
cropped_image = output_image[22:output_image.shape[0] - 40, 74:output_image.shape[1] - 33]
|
908 |
+
return cropped_image
|
909 |
+
|
910 |
+
if __name__ == "__main__":
|
911 |
+
|
912 |
+
image_size = [1024,1024]
|
913 |
+
target_distances = [650,1000,2000]
|
914 |
+
num_ys = 10
|
915 |
+
|
916 |
+
PATH_model_seg = 'C:/Users/valac/OneDrive - Západočeská univerzita v Plzni/DP/gradio/SegFormer_B3_1024_finetuned.pth'
|
917 |
+
PATH_model_det = 'C:/Users/valac/OneDrive - Západočeská univerzita v Plzni/DP/gradio/yolov8s.pt'
|
918 |
+
input_image = cv2.imread('C:/Users/valac/OneDrive - Západočeská univerzita v Plzni/DP/gradio/rs00006.jpg') #TO CO VLOZI UZIVATEL
|
919 |
+
model_seg = load_segformer(PATH_model_seg)
|
920 |
+
model_det = load_yolo(PATH_model_det)
|
921 |
+
image = run(input_image, model_seg, model_det, image_size, target_distances, num_ys=num_ys)
|