def test_subprocess_read(self):
shm = SharedMemory((100, 100), np.uint16)
a = np.random.RandomState(1234).randint(0, 65535, (100, 100))
with shm.txn() as m:
m[:] = a
with multiprocessing.Pool(1) as pool:
result = pool.apply(read_something, (shm, (55, 33)))
self.assertEqual(a[55, 33], result)
def test_subprocess_write(self):
shm = SharedMemory((100, 100), np.uint16)
with shm.txn() as m:
m[:] = 0
with multiprocessing.Pool(1) as pool:
pool.apply(write_something, (shm, (40, 50), 89))
with shm.txn() as m:
self.assertEqual(m[40, 50], 89)
def main():
args = parse_args()
source_files = sorted(glob.glob(args.source))
rb = RollingBuffer(source_files, args.n_io_cores)
points = np.array(json.load(open(args.points)))
patch_size = args.patch_size
half_patch_size = patch_size // 2
patches_xy, patches_xz, patches_yz = \
[SharedMemory((len(points), patch_size, patch_size), rb.dtype)
for _ in range(3)]
points_out = []
offset = 0
x1 = rb.shape[2] - half_patch_size
y1 = rb.shape[1] - half_patch_size
with multiprocessing.Pool(args.n_cores) as pool:
it = tqdm.tqdm(
range(half_patch_size, len(source_files) - half_patch_size))
for z in it:
it.set_description("Releasing @ %d" % (z - half_patch_size))
rb.release(z - half_patch_size)
it.set_description("Waiting for %d" % (z + half_patch_size))
rb.wait(z + half_patch_size)
pz = points[(points[:, 2] >= z) & (points[:, 2] < z+1)]
if len(pz) == 0:
continue
mask = np.all(pz[:, 0:2] >= half_patch_size, 1) &\
(pz[:, 0] < x1) &\
(pz[:, 1] < y1)
pz = pz[mask]
if len(pz) == 0:
continue
points_out.append(pz)
if len(pz) < args.n_cores * 10 or args.n_cores == 1:
offset = do_z(pz, offset,
patches_xy, patches_xz, patches_yz, rb, z,
half_patch_size)
else:
idxs = np.linspace(0, len(pz), args.n_cores+1).astype(int)
it.set_description("Freezing buffer")
frb = rb.freeze()
fnargs = [(pz[i0:i1], offset + i0,
patches_xy, patches_xz, patches_yz, frb, z,
half_patch_size)
for i0, i1 in zip(idxs[:-1], idxs[1:])]
it.set_description("Processing %d patches @ %d" % (len(pz), z))
pool.starmap(do_z, fnargs)
offset += len(pz)
points_out = np.vstack(points_out)
with h5py.File(args.output, "w") as f:
with patches_xy.txn() as m:
old_patches = f.create_dataset("patches_xy", data=m[:offset])
with patches_xz.txn() as m:
f.create_dataset("patches_xz", data=m[:offset])
with patches_yz.txn() as m:
f.create_dataset("patches_yz", data=m[:offset])
f.create_dataset("x", data=points_out[:, 0])
f.create_dataset("y", data=points_out[:, 1])
f.create_dataset("z", data=points_out[:, 2])
f["patches"] = old_patches
def main():
args = parse_arguments()
files = sorted(glob.glob(args.source))
first_plane = tifffile.imread(files[0])
x_extent = first_plane.shape[1]
y_extent = first_plane.shape[0]
z_extent = len(files)
x0a = np.arange(args.padding_xy, x_extent, args.block_size_xy)
x1a = x0a + args.block_size_xy
x1a[-1] = x_extent - args.padding_xy
x0p = x0a - args.padding_xy
x1p = x1a + args.padding_xy
x1p[-1] = x_extent
y0a = np.arange(args.padding_xy, y_extent, args.block_size_xy)
y1a = y0a + args.block_size_xy
y1a[-1] = y_extent - args.padding_xy
y0p = y0a - args.padding_xy
y1p = y1a + args.padding_xy
y1p[-1] = y_extent
points = []
for z0 in tqdm.trange(args.padding_z,
z_extent - args.padding_z,
args.block_size_z,
desc="Reading Z-block"):
z1 = min(z0 + args.block_size_z, z_extent - args.padding_z)
z0p = z0 - args.padding_z
z1p = min(z1 + args.padding_z, z_extent)
img_mem = SharedMemory((z1p - z0p, y_extent, x_extent),
first_plane.dtype)
with multiprocessing.Pool(12) as pool:
futures = []
for z in range(z0p, z1p):
futures.append(
pool.apply_async(read_plane, (img_mem, files[z], z - z0p)))
for future in tqdm.tqdm(futures, desc="Reading stack"):
future.get()
with multiprocessing.Pool() as pool:
futures = []
for xi, yi in itertools.product(range(len(x0a)), range(len(y0a))):
futures.append(
pool.apply_async(
do_dog,
(img_mem, args.dog_low, args.dog_high, x0a[xi],
x1a[xi], y0a[yi], y1a[yi], z0, z1, x0p[xi], x1p[xi],
y0p[yi], y1p[yi], z0p, z1p, args.min_distance,
args.threshold, args.invert)))
for future in tqdm.tqdm(futures, desc="Computing"):
points.append(future.get())
points = np.vstack(points)
with open(args.output, "w") as fd:
json.dump(points.tolist(), fd)
def do_plane(filename:str,
points:np.ndarray,
shared_memory:SharedMemory,
offset:int):
"""
:param filename: name of file to parse
:param points: an N x 2 array of X, Y points at which to sample
:param shared_memory: Shared memory block to write into
:param offset: offset into block
"""
patch_size = shared_memory.shape[1]
half_size = patch_size // 2
plane = np.pad(tifffile.imread(filename), half_size, mode='reflect')
for idx, (x, y) in enumerate(points):
x0 = x
x1 = x + patch_size
y0 = y
y1 = y + patch_size
with shared_memory.txn() as m:
m[offset + idx] = plane[y0:y1, x0:x1]
def write_something(shm: SharedMemory, idx, value):
with shm.txn() as m:
m[idx] = value
def read_something(shm: SharedMemory, idx):
with shm.txn() as m:
return m[idx]