def test_for_datatile_2(masks):
tile = DataTile(
tile_slice=Slice(origin=(0, 0, 0), shape=Shape((2 * 2, 10, 10), sig_dims=2)),
data=np.ones((2 * 2, 10, 10))
)
slice_ = masks.get_masks_for_slice(tile.tile_slice)
assert slice_.shape == (100, 5)
def get_tiles(self, tiling_scheme, dest_dtype="float32", roi=None):
yield from self._get_tiles_fullframe(tiling_scheme, dest_dtype, roi)
return
# assert len(tiling_scheme) == 1
print(tiling_scheme)
pool = self._data_source.pool.get_impl(
read_upto_frame=self._end_idx,
chunk_size=tiling_scheme.depth,
)
to_read = int(self._end_idx - self._start_idx)
nav_slices_raw = [(..., ) + slice_.get(sig_only=True)
for idx, slice_ in tiling_scheme.slices]
with pool:
while to_read > 0:
with pool.get_result() as res_wrapped:
frames_in_tile = res_wrapped.stop - res_wrapped.start
for (idx,
slice_), nav_slice_raw in zip(tiling_scheme.slices,
nav_slices_raw):
tile_shape = Shape(
(frames_in_tile, ) + tuple(slice_.shape),
sig_dims=2)
tile_slice = Slice(
origin=(res_wrapped.start, ) +
tuple(slice_.origin),
shape=tile_shape,
)
sliced_res = res_wrapped.buf[nav_slice_raw]
yield DataTile(sliced_res,
tile_slice=tile_slice,
scheme_idx=idx)
to_read -= frames_in_tile
def test_for_datatile_1(masks):
tile = DataTile(
tile_slice=Slice(origin=(0, 0, 0, 0), shape=(1, 1, 1, 1)),
data=np.ones((1, 1, 1, 1))
)
slice_ = masks.get_masks_for_slice(tile.tile_slice)
assert slice_.shape == (1, 4)
def _get_tiles_fullframe(self,
tiling_scheme,
dest_dtype="float32",
roi=None):
# assert len(tiling_scheme) == 1
logger.debug("reading up to frame idx %d for this partition",
self._end_idx)
pool = self._data_source.pool.get_impl(
read_upto_frame=self._end_idx,
# chunk_size=11,
chunk_size=tiling_scheme.depth,
)
to_read = self._end_idx - self._start_idx
with pool:
while to_read > 0:
# logger.debug(
# "LivePartition.get_tiles: to_read=%d, start_idx=%d end_idx=%d",
# to_read, self._start_idx, self._end_idx,
# )
with pool.get_result() as res_wrapped:
frames_in_tile = res_wrapped.stop - res_wrapped.start
tile_shape = Shape(
(frames_in_tile, ) + tuple(tiling_scheme[0].shape),
sig_dims=2)
tile_slice = Slice(
origin=(res_wrapped.start, ) + (0, 0),
shape=tile_shape,
)
yield DataTile(
res_wrapped.buf,
tile_slice=tile_slice,
scheme_idx=0,
)
to_read -= frames_in_tile
logger.debug("LivePartition.get_tiles: end of method")
def test_for_datatile_with_frame_origin(masks):
tile = DataTile(tile_slice=Slice(origin=(10, 10, 10, 10),
shape=(2, 2, 1, 5)),
data=np.ones((2, 2, 1, 5)))
slice_ = masks.get_masks_for_slice(tile.tile_slice)
print(slice_)
assert_array_almost_equal(
slice_,
np.array([
1,
0,
1,
10,
1,
0,
1,
11,
1,
0,
1,
12,
1,
0,
1,
13,
1,
0,
1,
14,
]).reshape((5, 4)))
def test_for_datatile_with_scan_origin(masks):
tile = DataTile(
tile_slice=Slice(origin=(10, 10, 0, 0), shape=(2, 2, 10, 10)),
data=np.ones((2, 2, 10, 10))
)
slice_ = masks.get_masks_for_slice(tile.tile_slice)
assert slice_.shape == (100, 4)
def test_for_datatile_with_frame_origin(masks):
tile = DataTile(tile_slice=Slice(origin=(10, 10, 10, 10),
shape=Shape((2, 2, 1, 5), sig_dims=2)),
data=np.ones((2, 2, 1, 5)))
slice_ = masks.get_masks_for_slice(tile.tile_slice)
print(slice_)
np.allclose(
slice_,
np.array([
1,
0,
1,
0,
10,
1,
0,
1,
0,
11,
1,
0,
1,
0,
12,
1,
0,
1,
0,
13,
1,
0,
1,
0,
14,
]).reshape((5, 5)))
def get_tiles(self):
subslices = list(self.slice.subslices(shape=self.tileshape))
for tile_slice in subslices:
yield DataTile(
data=self.dataset.data[tile_slice.get()],
tile_slice=tile_slice
)
def test_for_datatile_1(masks):
tile = DataTile(
np.ones((1, 1, 1)),
tile_slice=Slice(origin=(0, 0, 0), shape=Shape((1, 1, 1), sig_dims=2)),
scheme_idx=0,
)
slice_ = masks.get_masks_for_slice(tile.tile_slice)
assert slice_.shape == (1, 5)
def test_for_datatile_with_scan_origin(masks):
tile = DataTile(
np.ones((2 * 2, 10, 10)),
tile_slice=Slice(origin=(110, 0, 0),
shape=Shape((2 * 2, 10, 10), sig_dims=2)),
scheme_idx=0,
)
slice_ = masks.get_masks_for_slice(tile.tile_slice)
assert slice_.shape == (100, 5)
def get_tiles(self, *args, **kwargs):
crop = self._tileshape[1:]
for tile in super().get_tiles(*args, **kwargs):
log.debug("tile with slice %r" % tile.tile_slice)
for subslice in tile.tile_slice.subslices((self._stackheight,) + crop):
if self._tiledelay:
time.sleep(self._tiledelay)
yield DataTile(
data=subslice.shift(tile.tile_slice).get(tile.data),
tile_slice=subslice,
)
def get_tiles(self, crop_to=None):
stackheight = self.tileshape.nav.size
tile_buffer = np.zeros(self.tileshape.flatten_nav(),
dtype=self.dataset.dtype)
files_by_tile = list(grouper(self.files, stackheight))
sig_dims = self.tileshape.sig.dims
dtype = self.dataset.dtype
partition_origin = self.slice.origin
tile_slice_shape = Shape(tile_buffer.shape, sig_dims=sig_dims)
files_done = 0
for files in files_by_tile:
if len(files) != stackheight:
# allocate different size buffer for "rest-tile":
tile_buffer = np.zeros(
(len(files), ) + tuple(self.dataset.raw_shape.sig),
dtype=dtype)
tile_slice_shape = Shape(tile_buffer.shape, sig_dims=sig_dims)
origin = tuple(
sum(x) for x in zip(partition_origin, (files_done, 0, 0)))
tile_slice = Slice(
origin=origin,
shape=tile_slice_shape,
)
files_done += len(files)
if crop_to is not None:
intersection = tile_slice.intersection_with(crop_to)
if intersection.is_null():
continue
if stackheight == 1:
with open(files[0], "rb") as f:
f.readinto(tile_buffer)
yield DataTile(data=tile_buffer, tile_slice=tile_slice)
continue
else:
for idx, fn in enumerate(files):
with open(fn, "rb") as f:
f.readinto(tile_buffer[idx, ...])
yield DataTile(data=tile_buffer, tile_slice=tile_slice)
def get_tiles(self, crop_to=None, full_frames=False):
if full_frames:
tileshape = (tuple(self.tileshape[:self.meta.shape.nav.dims]) +
tuple(self.meta.shape.sig))
else:
tileshape = self.tileshape
subslices = self.slice.subslices(shape=tileshape)
for tile_slice in subslices:
if crop_to is not None:
intersection = tile_slice.intersection_with(crop_to)
if intersection.is_null():
continue
yield DataTile(data=self.reader.data[tile_slice.get()],
tile_slice=tile_slice)
def test_write_handle(tmpdir_factory):
"""
test the common "happy path":
"""
datadir = tmpdir_factory.mktemp('write_handle_tests')
full_path = os.path.join(datadir, "f1")
part_slice = Slice(
shape=Shape((32, 64, 64), sig_dims=2),
origin=(16, 0, 0),
)
tile_slice = Slice(
shape=Shape((3, 64, 64), sig_dims=2),
origin=(19, 0, 0),
)
tile_data = np.random.randn(3, 64, 64).astype("float32")
tile = DataTile(
tile_data,
tile_slice=tile_slice,
scheme_idx=0,
)
wh = WriteHandle(full_path, datadir, part_slice, dtype='float32')
tmp_file_name = ""
with wh:
wh.write_tile(tile)
tmp_file_name = wh._tmp_file.name
assert os.path.exists(tmp_file_name)
# check some internals:
assert wh._dest is None
assert wh._tmp_file is None
# the temporary file should no longer exist in case of success
assert not os.path.exists(tmp_file_name)
# ... buf our dest fname should:
assert os.path.exists(full_path)
assert os.path.isfile(full_path)
# check if data is written correctly:
read_data = np.fromfile(full_path,
dtype="float32").reshape(part_slice.shape)
assert np.allclose(read_data[3:6, ...], tile_data)
def _get_tiles_roi(self, tiling_scheme, fileset, read_ranges, roi, sync_offset):
ds_sig_shape = tiling_scheme.dataset_shape.sig
sig_dims = tiling_scheme.shape.sig.dims
slices, ranges, scheme_indices = read_ranges
fh = fileset[0]
memmap = fh.mmap().reshape((fh.num_frames,) + tuple(ds_sig_shape))
if sync_offset == 0:
flat_roi = roi.reshape((-1,))
elif sync_offset > 0:
flat_roi = np.full(roi.reshape((-1,)).shape, False)
flat_roi[:sync_offset] = False
flat_roi[sync_offset:] = roi.reshape((-1,))[:-sync_offset]
else:
flat_roi = np.full(roi.reshape((-1,)).shape, False)
flat_roi[sync_offset:] = False
flat_roi[:sync_offset] = roi.reshape((-1,))[-sync_offset:]
data_w_roi = memmap[flat_roi]
for idx in range(slices.shape[0]):
origin, shape = slices[idx]
scheme_idx = scheme_indices[idx]
tile_slice = Slice(
origin=origin,
shape=Shape(shape, sig_dims=sig_dims)
)
if sync_offset >= 0:
data_slice = tile_slice.get()
else:
frames_to_skip = np.count_nonzero(roi.reshape((-1,))[:abs(sync_offset)])
data_slice = Slice(
origin=(origin[0] - frames_to_skip,) + tuple(origin[-sig_dims:]),
shape=Shape(shape, sig_dims=sig_dims)
)
data_slice = data_slice.get()
data = data_w_roi[data_slice]
yield DataTile(
data,
tile_slice=tile_slice,
scheme_idx=scheme_idx,
)
def _get_tiles_roi(self, tiling_scheme, fileset, read_ranges, roi):
ds_sig_shape = tiling_scheme.dataset_shape.sig
sig_dims = tiling_scheme.shape.sig.dims
slices, ranges, scheme_indices = read_ranges
fh = fileset[0]
memmap = fh.mmap().reshape((fh.num_frames, ) + tuple(ds_sig_shape))
data_w_roi = memmap[roi.reshape((-1, ))]
for idx in range(slices.shape[0]):
origin, shape = slices[idx]
scheme_idx = scheme_indices[idx]
tile_slice = Slice(origin=origin,
shape=Shape(shape, sig_dims=sig_dims))
data_slice = tile_slice.get()
data = data_w_roi[data_slice]
yield DataTile(
data,
tile_slice=tile_slice,
scheme_idx=scheme_idx,
)
def test_write_handle_exception(tmpdir_factory):
datadir = tmpdir_factory.mktemp('write_handle_tests')
full_path = os.path.join(datadir, "f1")
part_slice = Slice(
shape=Shape((32, 64, 64), sig_dims=2),
origin=(16, 0, 0),
)
tile_slice = Slice(
shape=Shape((3, 64, 64), sig_dims=2),
origin=(19, 0, 0),
)
tile_data = np.random.randn(3, 64, 64).astype("float32")
tile = DataTile(
tile_data,
tile_slice=tile_slice,
scheme_idx=0,
)
wh = WriteHandle(full_path, datadir, part_slice, dtype='float32')
tmp_file_name = ""
with pytest.raises(Exception):
with wh:
wh.write_tile(tile)
tmp_file_name = wh._tmp_file.name
assert os.path.exists(tmp_file_name)
raise Exception("nope")
# check some internals:
assert wh._dest is None
assert wh._tmp_file is None
# the temporary file should no longer exist in case of exception
assert not os.path.exists(tmp_file_name)
# and neither should the full destination path
assert not os.path.exists(full_path)
