def write_to_kzip(trainer: 'Trainer', pred_batch: np.ndarray) -> None:
from knossos_utils import KnossosDataset
ks = trainer.knossos_preview_config
if isinstance(ks['dataset'], str):
dataset_path = ks['dataset']
else:
dataset_path = ks['dataset'][0]
ds = KnossosDataset(dataset_path)
seg = pred_batch[0].swapaxes(0, 2) # (N, D, H, W) -> (W, H, D)
# Set movement are in k.zip
area_min = ks['offset']
area_sz = ks['size']
anno_str = f"""<?xml version="1.0" encoding="UTF-8"?>
<things>
<parameters>
<MovementArea min.x="{area_min[0]}" min.y="{area_min[1]}" min.z="{area_min[2]}" size.x="{area_sz[0]}" size.y="{area_sz[1]}" size.z="{area_sz[2]}"/>
</parameters>
<comments/>
<branchpoints/>
</things>"""
kzip_path = f'{trainer.save_path}/preview_{trainer.step}.k.zip'
logger.info(f'Writing preview inference to {kzip_path}')
ds.save_to_kzip(data=seg,
data_mag=ks.get('mag', 1),
kzip_path=kzip_path,
offset=ks['offset'],
mags=ks.get('target_mags', [1, 2]),
gen_mergelist=False,
upsample=False,
fast_resampling=False,
annotation_str=anno_str)
def create_preview_batch_from_knossos(
knossos_preview_config: Dict[str, str]) -> torch.Tensor:
from knossos_utils import KnossosDataset
config = knossos_preview_config
required_keys = ['dataset', 'size', 'offset', 'mag', 'target_mags']
for k in required_keys:
if k not in config.keys():
raise ValueError(
f'Required key {k} missing from knossos_preview_config.')
logger.info(f'Loading preview region from dataset {config["dataset"]}')
datasets = config['dataset']
if isinstance(datasets, str):
datasets = [datasets]
inp_np = []
for idx, dataset_path in enumerate(datasets):
ds = KnossosDataset(dataset_path)
inp_np.append(
ds.load_raw(offset=config['offset'],
size=config['size'],
mag=config['mag'],
datatype=np.float32)) # (D, H, W)
inp_np = np.stack(inp_np, axis=0) # C, D, H, W
inp_np = inp_np[None] # (N, C, D, H, W)
inp_np = inp_np / config.get('scale_brightness', 1.)
inp = torch.from_numpy(inp_np)
return inp
def test_KnossosDataset_initialize_without_conf__boundary(tmp_path, boundary):
kd = KnossosDataset()
kd.initialize_without_conf(str(tmp_path),
boundary=boundary,
scale=(1, 1, 1),
experiment_name='test',
mags=[1],
verbose=True)
def test_KnossosDataset_initialize_without_conf__mags(tmp_path):
kd = KnossosDataset()
with pytest.raises(AssertionError):
kd.initialize_without_conf(str(tmp_path),
boundary=(7, 9, 10),
scale=(1, 1, 1),
experiment_name='test',
mags=None,
verbose=True)
def test_KnossosDataset_initialize_without_conf__conf_exist(tmp_path):
from pathlib import Path
kd = KnossosDataset()
kd.initialize_without_conf(str(tmp_path),
boundary=(7, 9, 10),
scale=(1, 1, 1),
experiment_name='test',
mags=[1],
verbose=True)
assert (Path(kd.conf_path).is_file())
def kd_factory(kd_path, channel='jpg'):
with warnings.catch_warnings():
warnings.simplefilter("ignore")
# knossos_utils expects a path to the knossos.conf file
if not kd_path.endswith('knossos.conf'):
kd_path += "/mag1/knossos.conf"
kd = KnossosDataset() # Sets initial values of object
# kd.set_channel(channel) # TODO: currently requires additional adjustment of the data type, i.e. setting the channel explicitely currently leads to uint32 <-> uint64 issues in the CS segmentation
# Initializes the dataset by parsing the knossos.conf in path + "mag1"
kd.initialize_from_knossos_path(kd_path)
return kd
def test_Knossosdataset__initalize_without_conf__robust_magfolder_detection(
tmp_path, existing_mag, expected_mag):
(tmp_path / existing_mag).mkdir()
kd = KnossosDataset()
kd.initialize_without_conf(str(tmp_path),
boundary=(7, 9, 10),
scale=(1, 1, 1),
experiment_name='test',
mags=[1],
verbose=True)
assert ((tmp_path / expected_mag).is_dir())
def test_KnossosDataset_initalize_without_conf__mags_1_make_mag_folder_False(
tmp_path):
kd = KnossosDataset()
with pytest.raises(AssertionError):
kd.initialize_without_conf(str(tmp_path),
boundary=(7, 9, 19),
scale=(1, 1, 1),
experiment_name='test',
mags=[1],
make_mag_folders=False,
verbose=True)
def create_preview_batch_from_knossos(knossos_preview_config: Dict[str, str]) -> torch.Tensor:
from knossos_utils import KnossosDataset
config = knossos_preview_config
required_keys = ['dataset', 'size', 'offset', 'mag', 'target_mags']
for k in required_keys:
if k not in config.keys():
raise ValueError(f'Required key {k} missing from knossos_preview_config.')
logger.info(f'Loading preview region from dataset {config["dataset"]}')
ds = KnossosDataset(config['dataset'])
inp_np = ds.load_raw(
offset=config['offset'],
size=config['size'],
mag=config['mag'],
datatype=np.float32
) # (D, H, W)
inp_np = inp_np[None, None] # -> (N, C, D, H, W)
inp_np = inp_np / config.get('scale_brightness', 1.)
inp = torch.from_numpy(inp_np)
return inp
def kd_factory(kd_path: str, channel: str = 'jpg'):
"""
Todo:
* Requires additional adjustment of the data type,
i.e. setting the channel explicitly currently leads to uint32 <->
uint64 issues in the CS segmentation.
Args:
kd_path: Path to the KnossosDataset.
channel: Channel which to use. Currently not used.
Returns:
"""
with warnings.catch_warnings():
warnings.simplefilter("ignore")
# knossos_utils expects a path to the knossos.conf file
if not kd_path.endswith('knossos.conf'):
kd_path += "/mag1/knossos.conf"
kd = KnossosDataset() # Sets initial values of object
# kd.set_channel(channel) #
# Initializes the dataset by parsing the knossos.conf in path + "mag1"
kd.initialize_from_knossos_path(kd_path)
return kd
def trafo_objects_to_kd(
realign_map,
obj_ids=None,
label_id=3,
kd_path='/wholebrain/scratch/areaxfs3/knossosdatasets/j0126_realigned_v4b_cbs_ext0_fix/',
path_folder='/wholebrain/scratch/areaxfs3/j0126_cset_paper/obj_mito_1037_3d_8'
):
with open(path_folder + '/direct_map.pkl', 'rb') as f:
obj_map = pkl.load(f)
kd = KnossosDataset()
kd.initialize_from_knossos_path(kd_path)
if obj_ids is None:
obj_ids = list(obj_map.keys())
time_start = time.time()
for i_obj_id, obj_id in enumerate(obj_ids):
voxels = np.array([], dtype=np.int32).reshape(0, 3)
for voxel_file in obj_map[obj_id]:
voxels = np.concatenate([
voxels,
np.load(path_folder + '/voxels/%d_%d.npz' %
(voxel_file[0], voxel_file[1]))["%d" % obj_id]
])
voxels_t, mean_shift = trafo_object(voxels, realign_map)
bb = np.array([np.min(voxels_t, axis=0),
np.max(voxels_t, axis=0)],
dtype=np.int)
if np.product(bb.shape) > 1e10:
print("Obj %d too large" % obj_id)
continue
voxels_t -= bb[0]
vx_cube = np.zeros(bb[1] - bb[0] + 1, dtype=np.uint64)
vx_cube[voxels_t[:, 0], voxels_t[:, 1], voxels_t[:, 2]] = label_id
kd.from_matrix_to_cubes(bb[0],
data=vx_cube,
as_raw=False,
overwrite=False)
dt = time.time() - time_start
print("Written obj %d at bb %s - time / object: %.3fs - eta: %.3fh" %
(obj_id, bb[0], dt / (i_obj_id + 1), dt /
(i_obj_id + 1) * len(obj_ids) / 3600))
from knossos_utils import KnossosDataset
from elektronn2.utils.plotting import scroll_plot, scroll_plot2
path = os.path.expanduser("~/lustre/sdorkenw/j0126_3d_rrbarrier/")
arr = KnossosArray(path, max_ram=1200, n_preload=2)
arr.preload([0, 0, 0], sync=True)
sh = np.array([120, 240, 240])
off = np.array([0, 0, 0]) + np.random.randint(0, 500, size=3)
# a = arr[0:80, 0:220, 0:230]
# a0 = arr[0:100, 0:400, 0:200]
# a1 = arr[0:100, 100:400+100, 100:200+100]
# b = arr[0:640, 0:640, 0:384]
# c = arr[500:600, 500:800, 500:800]
#
kds = KnossosDataset()
kds.initialize_from_knossos_path(path)
get = utils.timeit(kds.from_raw_cubes_to_matrix)
# a0 = arr.cut_slice(sh, off)
# a1 = get(sh[[2,1,0]], off[[2,1,0]])
# a2 = (a1.T).astype(np.float32)/255
# print(np.allclose(a0, a2))
# fig = scroll_plot2([a0, a2], 'ab')
t = 0
out = np.empty(sh, dtype=np.float32)
N = 500
for i in range(N):
t0 = time.time()
a = arr.cut_slice(sh, off, out=out)