def test_upsample_multi_channel(tmp_path: Path) -> None:
num_channels = 3
size = (32, 32, 10)
source_data = (
128 * np.random.randn(num_channels, size[0], size[1], size[2])
).astype("uint8")
ds = Dataset(tmp_path / "multi-channel-test", (1, 1, 1))
l = ds.add_layer(
"color",
COLOR_CATEGORY,
dtype_per_channel="uint8",
num_channels=num_channels,
)
mag2 = l.add_mag("2", chunks_per_shard=32)
mag2.write(source_data)
assert np.any(source_data != 0)
l._initialize_mag_from_other_mag("1", mag2, False)
upsample_cube_job(
(mag2.get_view(), l.get_mag("1").get_view(), 0),
[0.5, 0.5, 0.5],
BUFFER_SHAPE,
)
channels = []
for channel_index in range(num_channels):
channels.append(upsample_cube(source_data[channel_index], [2, 2, 2]))
joined_buffer = np.stack(channels)
target_buffer = l.get_mag("1").read()
assert np.any(target_buffer != 0)
assert np.all(target_buffer == joined_buffer)
def test_default_anisotropic_voxel_size(tmp_path: Path) -> None:
ds = Dataset(tmp_path / "default_anisotropic_voxel_size",
voxel_size=(85, 85, 346))
layer = ds.add_layer("color", COLOR_CATEGORY)
mag = layer.add_mag(1)
mag.write(data=(np.random.rand(10, 20, 30) * 255).astype(np.uint8))
layer.downsample(Mag(1), None, "median", True)
assert sorted(layer.mags.keys()) == [Mag("1"), Mag("2-2-1"), Mag("4-4-1")]
def test_downsample_with_invalid_mag_list(tmp_path: Path) -> None:
ds = Dataset(tmp_path / "downsample_mag_list", voxel_size=(1, 1, 2))
layer = ds.add_layer("color", COLOR_CATEGORY)
mag = layer.add_mag(1)
mag.write(data=(np.random.rand(10, 20, 30) * 255).astype(np.uint8))
with pytest.raises(AssertionError):
layer.downsample_mag_list(
from_mag=Mag(1),
target_mags=[Mag(1),
Mag([1, 1, 2]),
Mag([2, 2, 1]),
Mag(2)],
)
def test_downsample_mag_list(tmp_path: Path) -> None:
ds = Dataset(tmp_path / "downsample_mag_list", voxel_size=(1, 1, 2))
layer = ds.add_layer("color", COLOR_CATEGORY)
mag = layer.add_mag(1)
mag.write(data=(np.random.rand(10, 20, 30) * 255).astype(np.uint8))
target_mags = [Mag([4, 4, 8]),
Mag(2), Mag([32, 32, 8]),
Mag(32)] # unsorted list
layer.downsample_mag_list(from_mag=Mag(1), target_mags=target_mags)
for m in target_mags:
assert m in layer.mags
def test_default_parameter(tmp_path: Path) -> None:
target_path = tmp_path / "downsaple_default"
ds = Dataset(target_path, voxel_size=(1, 1, 1))
layer = ds.add_layer("color",
COLOR_CATEGORY,
dtype_per_channel="uint8",
num_channels=3)
mag = layer.add_mag("2")
mag.write(data=(np.random.rand(3, 10, 20, 30) * 255).astype(np.uint8))
layer.downsample()
# The max_mag is Mag(4) in this case (see test_default_max_mag)
assert sorted(layer.mags.keys()) == [Mag("2"), Mag("4")]
def test_upsampling(tmp_path: Path) -> None:
ds = Dataset(tmp_path, voxel_size=(1, 1, 1))
layer = ds.add_layer("color", COLOR_CATEGORY)
mag = layer.add_mag([4, 4, 2])
mag.write(
absolute_offset=(10 * 4, 20 * 4, 40 * 2),
data=(np.random.rand(46, 45, 27) * 255).astype(np.uint8),
)
layer.upsample(
from_mag=Mag([4, 4, 2]),
finest_mag=Mag(1),
compress=False,
sampling_mode=SamplingModes.ANISOTROPIC,
buffer_edge_len=64,
args=None,
)
def test_downsample_2d(tmp_path: Path) -> None:
ds = Dataset(tmp_path / "downsample_compressed", voxel_size=(1, 1, 2))
layer = ds.add_layer("color", COLOR_CATEGORY)
mag = layer.add_mag(1, chunk_size=8, chunks_per_shard=8)
# write 2D data with all values set to "123"
mag.write(data=(np.ones((100, 100, 1)) * 123).astype(np.uint8))
with pytest.warns(Warning):
# This call produces a warning because only the mode "CONSTANT_Z" makes sense for 2D data.
layer.downsample(
from_mag=Mag(1),
coarsest_mag=Mag(2),
sampling_mode=SamplingModes.
ISOTROPIC, # this mode is intentionally not "CONSTANT_Z" for this test
)
assert Mag("2-2-1") in layer.mags
assert np.all(layer.get_mag(
Mag("2-2-1")).read() == 123) # The data is not darkened
def test_downsample_multi_channel(tmp_path: Path) -> None:
num_channels = 3
size = (32, 32, 10)
source_data = (128 * np.random.randn(num_channels, size[0], size[1],
size[2])).astype("uint8")
ds = Dataset(tmp_path / "multi-channel-test", (1, 1, 1))
l = ds.add_layer(
"color",
COLOR_CATEGORY,
dtype_per_channel="uint8",
num_channels=num_channels,
)
mag1 = l.add_mag("1", chunks_per_shard=32)
print("writing source_data shape", source_data.shape)
mag1.write(source_data)
assert np.any(source_data != 0)
mag2 = l._initialize_mag_from_other_mag("2", mag1, False)
downsample_cube_job(
(l.get_mag("1").get_view(), l.get_mag("2").get_view(), 0),
Vec3Int(2, 2, 2),
InterpolationModes.MAX,
BUFFER_SHAPE,
)
channels = []
for channel_index in range(num_channels):
channels.append(
downsample_cube(source_data[channel_index], [2, 2, 2],
InterpolationModes.MAX))
joined_buffer = np.stack(channels)
target_buffer = mag2.read()
assert np.any(target_buffer != 0)
assert np.all(target_buffer == joined_buffer)
def upsample_test_helper(tmp_path: Path, use_compress: bool) -> None:
ds = Dataset(tmp_path, voxel_size=(10.5, 10.5, 5))
layer = ds.add_layer("color", COLOR_CATEGORY)
mag2 = layer.add_mag([2, 2, 2])
offset = Vec3Int(WKW_CUBE_SIZE, 2 * WKW_CUBE_SIZE, 0)
mag2.write(
absolute_offset=offset,
data=(np.random.rand(*BUFFER_SHAPE) * 255).astype(np.uint8),
)
mag1 = layer._initialize_mag_from_other_mag("1-1-2", mag2, use_compress)
source_buffer = mag2.read(
absolute_offset=offset,
size=BUFFER_SHAPE,
)[0]
assert np.any(source_buffer != 0)
upsample_cube_job(
(
mag2.get_view(absolute_offset=offset, size=BUFFER_SHAPE),
mag1.get_view(
absolute_offset=offset,
size=BUFFER_SHAPE,
),
0,
),
[0.5, 0.5, 1.0],
BUFFER_SHAPE,
)
assert np.any(source_buffer != 0)
target_buffer = mag1.read(absolute_offset=offset, size=BUFFER_SHAPE)[0]
assert np.any(target_buffer != 0)
assert np.all(target_buffer == upsample_cube(source_buffer, [2, 2, 1]))
def test_downsample_compressed(tmp_path: Path) -> None:
ds = Dataset(tmp_path / "downsample_compressed", voxel_size=(1, 1, 2))
layer = ds.add_layer("color", COLOR_CATEGORY)
mag = layer.add_mag(1, chunk_size=8, chunks_per_shard=8)
mag.write(data=(np.random.rand(80, 240, 15) * 255).astype(np.uint8))
assert not mag._is_compressed()
mag.compress()
assert mag._is_compressed()
layer.downsample(
from_mag=Mag(1),
coarsest_mag=Mag(
4
), # Setting max_mag to "4" covers an edge case because the z-dimension (15) has to be rounded
)
# Note: this test does not check if the data is correct. This is already covered by other test cases.
assert len(layer.mags) == 3
assert Mag("1") in layer.mags.keys()
assert Mag("2-2-1") in layer.mags.keys()
assert Mag("4-4-2") in layer.mags.keys()
def test_downsample_mag_list_with_only_setup_mags(tmp_path: Path) -> None:
ds = Dataset(tmp_path / "downsample_mag_list", voxel_size=(1, 1, 2))
layer = ds.add_layer("color", COLOR_CATEGORY)
mag = layer.add_mag(1)
mag.write(data=(np.random.rand(10, 20, 30) * 255).astype(np.uint8))
target_mags = [Mag([4, 4, 8]),
Mag(2), Mag([32, 32, 8]),
Mag(32)] # unsorted list
layer.downsample_mag_list(from_mag=Mag(1),
target_mags=target_mags,
only_setup_mags=True)
for m in target_mags:
assert np.all(
layer.get_mag(m).read() == 0), "The mags should be empty."
layer.downsample_mag_list(from_mag=Mag(1),
target_mags=target_mags,
allow_overwrite=True)
for m in target_mags:
assert m in layer.mags