oValach commited on
Commit
a0c9527
1 Parent(s): 3a47436

Delete TheDistanceAssessor.py

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