forked from ant-trullo/SegmentTrack_v4.0
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathNucleiConnectMultiCore.py
52 lines (39 loc) · 2.34 KB
/
NucleiConnectMultiCore.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
"""This function trackes nuclei for very big data.
It just coordinates the previous function 'NucleiConnect' to work in a
multiprocessing pool.
"""
import multiprocessing
from importlib import reload
import numpy as np
import NucleiConnect
class NucleiConnectMultiCore:
"""Main class, does all the job"""
def __init__(self, nuclei_seg, dist_thr):
reload(NucleiConnect)
steps = nuclei_seg.shape[0]
cpu_ow = multiprocessing.cpu_count()
chops = np.round(np.linspace(0, steps, cpu_ow + 1)).astype(np.int)
if steps > 10 * cpu_ow:
job_args = []
for t in range(chops.size - 1):
job_args.append([nuclei_seg[chops[t]:chops[t + 1], :, :], dist_thr])
pool = multiprocessing.Pool()
results = pool.map(NucleiConnect.NucleiConnect, job_args)
pool.close()
nuclei_tracked = np.zeros(nuclei_seg.shape, dtype=np.int32)
nuclei_tracked[chops[0]:chops[1], :, :] = results[0].nuclei_tracked
for t in range(1, cpu_ow):
idx = np.unique(results[t].nuclei_tracked[0, :, :])[1:] # after pooling, results must be concatenate but saving the correct tag for each nucleus:
for k in idx: # here we work at the interface (last frame versus first frame of the following results block)
tag = (nuclei_tracked[chops[t] - 1, :, :] * (results[t].nuclei_tracked[0, :, :] == k))
tag = tag.reshape(tag.size)
tag = np.delete(tag, np.where(tag == 0))
if tag.size > 0:
tag = np.median(tag).astype(np.int32)
nuclei_tracked[chops[t]:chops[t + 1], :, :] += (results[t].nuclei_tracked == k) * tag
else:
nuclei_tracked = NucleiConnect.NucleiConnect([nuclei_seg, dist_thr]).nuclei_tracked
# mycmap = np.fromfile("mycmap.bin", "uint16").reshape((10000, 3)) / 255.0
# nuclei_tracked_visual = label2rgb(nuclei_tracked, bg_label=0, bg_color=[0, 0, 0], colors=mycmap)
self.nuclei_tracked = nuclei_tracked
# self.nuclei_tracked_visual = nuclei_tracked_visual