-
Notifications
You must be signed in to change notification settings - Fork 0
/
dbscan.py
274 lines (207 loc) · 9.55 KB
/
dbscan.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
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
import argparse
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.cm import ScalarMappable
from matplotlib.colors import ListedColormap
from sklearn.cluster import DBSCAN
# see https://scikit-learn.org/stable/auto_examples/cluster/plot_dbscan.html
def main():
coordinates_3d = np.load(TRACK_3D_NPY_FILE_PATH).astype(float)
coordinates_3d = convert_to_micro_meters(coordinates_3d)
time = coordinates_3d[:, 3]
coordinates_3d_with_time = coordinates_3d[:, :4]
# -- Plot before dbscan to see how everything looks like --------------------------------------
axis_3d_no_dbscan = plot_3d_with_color(coordinates_3d_with_time)
# plot_color_bar_of_timepoints(time)
# plt.savefig(f'{TRACK_3D_NPY_BASE_FILENAME}.png', format='png')
# plt.show()
# -- ------------------------------------------------------------------------------------------
## -- this is all it takes for DBSCAN ---------------------------------------------------------
normalized_centers_with_time = normalize(coordinates_3d_with_time)
db = DBSCAN(eps=DBSCAN_EPS, min_samples=DBSCAN_MIN_SAMPLES, metric=custom_distance)
db.fit(normalized_centers_with_time)
# -- ------------------------------------------------------------------------------------------
n_clusters_ = get_num_of_clusters(db)
n_noise_ = get_num_of_noise_points(db)
log_n_clusters_and_n_noise(n_clusters_, n_noise_)
# -- Plot after dbscan - no clouds around clusters - noise points removed ---------------------
# axis_3d_no_noise = plot_3d_without_noise_no_cloud(coordinates_3d_with_time, db)
# plt.show()
# -- ------------------------------------------------------------------------------------------
# -- Plot after dbscan - plot the clouds around clusters and save results to disk--------------
plot_3d_with_clouds(coordinates_3d,
db,
plot_noise=PLOT_NOISE_POINTS,
save_clusters=SAVE_CLUSTERS_TO_DISK,
axis_3d=axis_3d_no_dbscan)
if PLOT_NOISE_POINTS:
plt.savefig(f'{TRACK_3D_NPY_BASE_FILENAME}_clusters_{n_clusters_}_n_noise_points_{n_noise_}.png', format='png')
else:
plt.savefig(
f'{TRACK_3D_NPY_BASE_FILENAME}_clusters_{n_clusters_}_n_noise_points_{n_noise_}_no_noise_points.png',
format='png')
plt.show()
# -- ------------------------------------------------------------------------------------------
def normalize(coordinates_3d_with_time):
# calculate the norm of each column
max_values = np.max(coordinates_3d_with_time, axis=0)
# normalize the columns
normalized_centers_with_time = coordinates_3d_with_time / max_values
return normalized_centers_with_time
def weighted_point(p):
wp = np.array(p[0:4])
# scale z axis
wp[2] = 4 * wp[2]
return wp
def custom_distance(p1, p2):
dist = np.linalg.norm(weighted_point(p1) - weighted_point(p2))
return dist
def convert_to_micro_meters(coordinates_3d):
# convert pixels to micro meter
coordinates_3d[:, 0:2] = coordinates_3d[:, 0:2] * (330.0 / 2048) # convert X-Y from pixel to micro meter values
coordinates_3d[:, 2] = coordinates_3d[:, 2] * 4 # convert z to correct micro meter values
return coordinates_3d
def create_time_points_colors(time):
time_colors = plt.cm.rainbow(np.linspace(0, 1, int(max(time)) + 1))
time_points_colors = [time_colors[int(tp)] for tp in time]
return time_points_colors
def get_num_of_noise_points(db):
n_noise_ = list(db.labels_).count(-1)
return n_noise_
def get_num_of_clusters(db):
labels = db.labels_
unique_labels = set(labels)
# number of clusters in labels, ignoring noise if present.
n_clusters_ = len(unique_labels) - (1 if -1 in labels else 0)
return n_clusters_
def create_core_samples_mask(db):
labels = db.labels_
core_samples_mask = np.zeros_like(labels, dtype=bool)
core_samples_mask[db.core_sample_indices_] = True
return core_samples_mask
def setup_3d_plot():
# setup 3d plot
fig = plt.figure(figsize=(12, 8))
axis = fig.add_subplot(111, projection='3d')
axis.set_xlim(0, 365)
axis.set_ylim(0, 365)
axis.set_zlim(0, 400)
axis.set_xlabel("x [µm]")
axis.set_ylabel("y [µm]")
axis.set_zlabel("z [µm]")
axis.view_init(elev=40, azim=-166)
return axis
def plot_3d_with_color(coordinates_3d_with_time):
time = coordinates_3d_with_time[:, 3]
time_points_colors = create_time_points_colors(time)
axis_3d = setup_3d_plot()
axis_3d.scatter(coordinates_3d_with_time[:, 0]
, coordinates_3d_with_time[:, 1]
, coordinates_3d_with_time[:, 2]
, "o"
, s=10
, facecolor=time_points_colors
, edgecolors=time_points_colors)
return axis_3d
def plot_color_bar_of_timepoints(time):
time_points_colors = create_time_points_colors(time)
custom_cmap = ListedColormap(time_points_colors)
# create a custom Normalize object to map colors to a range of values
norm = mpl.colors.Normalize(vmin=0, vmax=len(time_points_colors) - 1)
# create a ScalarMappable object with the custom colormap and normalization
sm = ScalarMappable(cmap=custom_cmap, norm=norm)
sm.set_array([])
# create a color bar using the ScalarMappable object
cbar = plt.colorbar(sm, orientation='horizontal')
cbar.set_label('Time index', fontsize=14)
cbar.ax.tick_params(labelsize=12)
def plot_3d_without_noise_no_cloud(coordinates_3d_with_time, db):
core_samples_mask = create_core_samples_mask(db)
labels = db.labels_
axis_3d = setup_3d_plot()
mask = labels == -1
xyzt = coordinates_3d_with_time[~mask & core_samples_mask]
plot_3d_with_color(xyzt)
return axis_3d
def plot_3d_with_clouds(coordinates_3d, db, plot_noise, save_clusters, axis_3d=None):
if axis_3d is None:
axis_3d = setup_3d_plot()
coordinates_3d_with_time = coordinates_3d[:, :4]
labels = db.labels_
unique_labels = set(labels)
n_clusters_ = get_num_of_clusters(db)
n_noise_ = get_num_of_noise_points(db)
label_colors = [plt.cm.rainbow(each) for each in np.linspace(0, 1, len(unique_labels))]
core_samples_mask = create_core_samples_mask(db)
plt.title(f'Estimated number of clusters: {n_clusters_}\n'
f'Estimated number of noise points: {n_noise_}')
# create a legend with label colors
handles = [plt.Rectangle((0, 0), 1, 1, fc=tuple(label_colors[i])) for i in range(len(unique_labels))]
legend_labels = [f'Track {i + 1}' for i in range(len(unique_labels) - 1)]
axis_3d.legend(handles, legend_labels, framealpha=0.0)
for k, col in zip(unique_labels, label_colors):
# use black for noise points
if k == -1:
col = [0, 0, 0, 1]
if not plot_noise:
continue
class_member_mask = labels == k
# -- comment in if you want to see the class number on each point
'''
xyzt = coordinates_3d_with_time[class_member_mask & core_samples_mask]
x = xyzt[:, 0]
y = xyzt[:, 1]
z = xyzt[:, 2]
for i, txt in enumerate(labels[class_member_mask & core_samples_mask]):
axis_3d.text(x[i], y[i], z[i], str(txt + 1), color='black', fontsize=8)
'''
xyzt = coordinates_3d_with_time[class_member_mask & core_samples_mask]
axis_3d.scatter(xyzt[:, 0]
, xyzt[:, 1]
, xyzt[:, 2]
, facecolor=tuple(col)
, s=250
, alpha=0.25)
xyzt = coordinates_3d_with_time[class_member_mask & ~core_samples_mask]
axis_3d.scatter(xyzt[:, 0]
, xyzt[:, 1]
, xyzt[:, 2]
, "o"
, facecolor=tuple(col)
, s=60
, alpha=0.25)
if save_clusters:
xyzt_and_diameter = coordinates_3d[class_member_mask]
if k == -1:
np.save(f'{TRACK_3D_NPY_BASE_FILENAME}_cluster_{k}_noise_points', xyzt_and_diameter)
else:
np.save(f'{TRACK_3D_NPY_BASE_FILENAME}_cluster_{k}', xyzt_and_diameter)
def log_n_clusters_and_n_noise(n_clusters_, n_noise_):
print("Estimated number of clusters: %d" % n_clusters_)
print("Estimated number of noise points: %d" % n_noise_)
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description='This script performs DBSCAN on the provided 3D track and plots the results.')
parser.add_argument('--track-3d', '-ts3d',
type=str,
help='path to the 3D track segment npy-file',
default='./')
parser.add_argument('--dbscan-eps', '-eps',
type=float,
help='DBSCAN EPS parameter',
default=0.15)
parser.add_argument('--dbscan-min-samples', '-ms',
type=int,
help='DBSCAN min_samples parameter',
default=10)
args = parser.parse_args()
# -- Change parameters here to avoid running the script from CLI if needed
TRACK_3D_NPY_FILE_PATH = args.track_3d
DBSCAN_EPS = args.dbscan_eps
DBSCAN_MIN_SAMPLES = args.dbscan_min_samples
# --------------------------------------------------------
TRACK_3D_NPY_BASE_FILENAME = TRACK_3D_NPY_FILE_PATH[:-4]
PLOT_NOISE_POINTS = False
SAVE_CLUSTERS_TO_DISK = False
main()