This repository has been archived by the owner on Apr 1, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathregion_growing_4.py
151 lines (133 loc) · 5.44 KB
/
region_growing_4.py
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
import os
import pickle
from datetime import datetime
from typing import Optional
import cv2
import matplotlib.pyplot as plot
import numpy as np
from config import bitmap, bitmap_folder, dim, min_seg
# read the image
loading_time = datetime.now()
arr: np.ndarray = cv2.cvtColor(cv2.imread(os.path.join('vis', 'output', bitmap_folder, bitmap + '.bmp')),
cv2.COLOR_BGR2YUV)
# YCbCr (correspondant to YUV) != YCrCb (provided by `cv2`)
print('Loading time:', datetime.now() - loading_time)
class Segment:
def __init__(self):
self.id: int = len(segments) + 1
self.p: list[int] = [] # pixels
self.ys: int = 0 # sum of Y channel
self.us: int = 0 # sum of U channel
self.vs: int = 0 # sum of V channel
self.m: list[int] = [] # average colour
def calculate_mean(self, al: int):
self.m = [round(self.ys / al), round(self.us / al), round(self.vs / al)]
def compare_colours(a: np.ndarray, b: np.ndarray) -> bool:
return abs(int(a[0]) - int(b[0])) <= 4 \
and abs(int(a[1]) - int(b[1])) <= 4 \
and abs(int(a[2]) - int(b[2])) <= 4
def find_a_segment_to_dissolve_in(seg_: Segment) -> Optional[tuple[int, int]]:
y_, x_ = seg_.p[0] >> 16, seg_.p[0] & 0xFFFF
if y_ > 0:
return y_ - 1, x_
if x_ > 0:
return y_, x_ - 1
last = len(seg_.p) - 1
y_, x_ = seg_.p[last] >> 16, seg_.p[last] & 0xFFFF
if y_ < dim - 1:
return y_ + 1, x_
if x_ < dim - 1:
return y_, x_ + 1
return None
# detect segments by single-pixel colour differences
segmentation_time = datetime.now()
status: np.ndarray = np.zeros((dim, dim), np.int32)
segments: list[Segment] = []
stack: list[list[int, int, int]] = []
thisY, thisX, found_sth_to_analyse = 0, 0, True
while found_sth_to_analyse:
found_sth_to_analyse = False
for y in range(thisY, dim):
for x in range(thisX if y == thisY else 0, dim):
if status[y, x] == 0:
found_sth_to_analyse = True
thisY = y
thisX = x
break
if found_sth_to_analyse: break
if not found_sth_to_analyse: break
seg = Segment()
stack.append([thisY, thisX, 0])
while len(stack) != 0:
l_ = len(stack) - 1
yy, xx, dr = stack[l_]
if dr == 0:
seg.p.append((yy << 16) | xx)
status[yy, xx] = seg.id
# if dr <= 0: # left
stack[l_][2] += 1
if xx > 0 and status[yy, xx - 1] == 0 and compare_colours(arr[yy, xx], arr[yy, xx - 1]):
stack.append([yy, xx - 1, 0])
continue
if dr <= 1: # top
stack[l_][2] += 1
if yy > 0 and status[yy - 1, xx] == 0 and compare_colours(arr[yy, xx], arr[yy - 1, xx]):
stack.append([yy - 1, xx, 0])
continue
if dr <= 2: # right
stack[l_][2] += 1
if xx < (dim - 1) and status[yy, xx + 1] == 0 and compare_colours(arr[yy, xx], arr[yy, xx + 1]):
stack.append([yy, xx + 1, 0])
continue
if dr <= 3: # bottom
stack[l_][2] += 1
if yy < (dim - 1) and status[yy + 1, xx] == 0 and compare_colours(arr[yy, xx], arr[yy + 1, xx]):
stack.append([yy + 1, xx, 0])
continue
stack.pop()
segments.append(seg)
# dissolve smaller segments
if min_seg > 1:
dissolution_time = datetime.now()
for seg in range(len(segments) - 1, -1, -1):
if len(segments[seg].p) < min_seg:
absorber_index = find_a_segment_to_dissolve_in(segments[seg])
if absorber_index is None: continue # rarely
absorber: Segment = segments[status[*absorber_index] - 1] # `- 1` is because segment IDs start from 0!
for p in segments[seg].p:
absorber.p.append(p)
status[p >> 16, p & 0xFFFF] = absorber.id
segments.pop(seg)
print('Dissolution time:', datetime.now() - dissolution_time)
print('+ Segmentation time:', datetime.now() - segmentation_time)
# it is actually a temporary part of tracing, but is done here for efficiency.
average_colours_time = datetime.now()
for seg in segments:
for p in seg.p:
y, x = p >> 16, p & 0xFFFF
seg.ys += arr[y, x][0]
seg.us += arr[y, x][1]
seg.vs += arr[y, x][2]
seg.calculate_mean(len(seg.p))
print('+ Average colours time:', datetime.now() - average_colours_time)
# evaluate the segments
segments.sort(key=lambda s: len(s.p), reverse=True)
arr = cv2.cvtColor(cv2.cvtColor(arr, cv2.COLOR_YUV2RGB), cv2.COLOR_RGB2HSV)
for big_sgm in range(25): # colour the biggest ones
for px in segments[big_sgm].p:
arr[px >> 16, px & 0xFFFF] = 5 + (10 * (big_sgm + 1)), 255, 255
for seg in range(len(segments)): # show the persisting small segments
if len(segments[seg].p) < min_seg:
for px in segments[seg].p:
arr[px >> 16, px & 0xFFFF] = 0, 255, 255
continue
arr = cv2.cvtColor(cv2.cvtColor(arr, cv2.COLOR_HSV2RGB), cv2.COLOR_RGB2YUV)
print('Total segments:', len(segments))
# show the image
plot.imshow(cv2.cvtColor(arr, cv2.COLOR_YUV2RGB))
plot.show()
# save the output
dumping_time = datetime.now()
pickle.dump(status, open(os.path.join('segmentation', 'output', 'rg4_' + bitmap + '_status.pickle'), 'wb'))
pickle.dump(segments, open(os.path.join('segmentation', 'output', 'rg4_' + bitmap + '_segments.pickle'), 'wb'))
print('Dumping time:', datetime.now() - dumping_time)