def __read_spec(self, array_key, hdf_file, ds_name):
dataset = hdf_file[ds_name]
dims = Coordinate(dataset.shape)
if self.ndims is None:
self.ndims = len(dims)
else:
assert self.ndims == len(dims)
if array_key in self.array_specs:
spec = self.array_specs[array_key].copy()
else:
spec = ArraySpec()
if spec.voxel_size is None:
if 'resolution' in dataset.attrs:
spec.voxel_size = Coordinate(dataset.attrs['resolution'])
else:
spec.voxel_size = Coordinate((1, ) * self.ndims)
logger.warning(
"WARNING: File %s does not contain resolution information "
"for %s (dataset %s), voxel size has been set to %s. This "
"might not be what you want.", self.filename, array_key,
ds_name, spec.voxel_size)
if spec.roi is None:
if 'offset' in dataset.attrs:
offset = Coordinate(dataset.attrs['offset'])
else:
offset = Coordinate((0, ) * self.ndims)
spec.roi = Roi(offset, dims * spec.voxel_size)
if spec.dtype is not None:
assert spec.dtype == dataset.dtype, (
"dtype %s provided in array_specs for %s, "
"but differs from dataset %s dtype %s" %
(self.array_specs[array_key].dtype, array_key, ds_name,
dataset.dtype))
else:
spec.dtype = dataset.dtype
if spec.interpolatable is None:
spec.interpolatable = spec.dtype in [
np.float,
np.float32,
np.float64,
np.float128,
np.uint8 # assuming this is not used for labels
]
logger.warning(
"WARNING: You didn't set 'interpolatable' for %s "
"(dataset %s). Based on the dtype %s, it has been "
"set to %s. This might not be what you want.", array_key,
ds_name, spec.dtype, spec.interpolatable)
return spec
def __init__(self, array: Array, key: gp.ArrayKey):
self.array = array
self.array_spec = ArraySpec(roi=self.array.roi,
voxel_size=self.array.voxel_size)
self.key = key
def __read_spec(self, headers):
num_files = len(headers)
assert num_files > 0
common_header = headers[0]
for header in headers:
for attr in ['imagesize_tczyx', 'pixelspacing_tczyx']:
assert (common_header[attr] == header[attr]).all(), (
"Headers of provided KLB files differ in attribute %s"%attr)
assert common_header['datatype'] == header['datatype'], (
"Headers of provided KLB files differ in attribute datatype")
size = Coordinate(common_header['imagesize_tczyx'])
voxel_size = Coordinate(common_header['pixelspacing_tczyx'])
dtype = common_header['datatype']
# strip leading 1 dimensions
while size[0] == 1 and len(size) > 1:
size = size[1:]
voxel_size = voxel_size[1:]
# append num_files dimension
if num_files > 1:
size = (num_files,) + size
voxel_size = (1,) + voxel_size
dims = Coordinate(size)
self.ndims = len(dims)
if self.array_spec is not None:
spec = self.array_spec
else:
spec = ArraySpec()
if spec.voxel_size is None:
spec.voxel_size = Coordinate(voxel_size)
if spec.roi is None:
offset = Coordinate((0,)*self.ndims)
spec.roi = Roi(offset, dims*spec.voxel_size)
if spec.dtype is not None:
assert spec.dtype == dtype, (
"dtype %s provided in array_specs for %s, but differs from "
"dataset dtype %s"%(
self.array_specs[self.array].dtype, self.array,
dataset.dtype))
else:
spec.dtype = dtype
if spec.interpolatable is None:
spec.interpolatable = spec.dtype in [
np.float,
np.float32,
np.float64,
np.float128,
np.uint8 # assuming this is not used for labels
]
logger.warning("WARNING: You didn't set 'interpolatable' for %s. "
"Based on the dtype %s, it has been set to %s. "
"This might not be what you want.",
self.array, spec.dtype, spec.interpolatable)
return spec
def add_embedding_training(
pipeline,
setup_config,
raw,
gt_labels,
mask,
neighborhood_gt=None,
neighborhood_mask=None,
):
# Network params
voxel_size = Coordinate(setup_config["VOXEL_SIZE"])
input_size = Coordinate(setup_config["INPUT_SHAPE"]) * voxel_size
output_size = Coordinate(setup_config["OUTPUT_SHAPE"]) * voxel_size
# Config options
embedding_net_name = setup_config["EMBEDDING_NET_NAME"]
checkpoint_every = setup_config["CHECKPOINT_EVERY"]
tensorboard_log_dir = setup_config["TENSORBOARD_LOG_DIR"]
# New array Keys
embedding = ArrayKey("EMBEDDING")
embedding_gradient = ArrayKey("EMBEDDING_GRADIENT")
neighborhood = ArrayKey("NEIGHBORHOOD")
neighborhood_gradient = ArrayKey("NEIGHBORHOOD_GRADIENT")
# model, optimizer, loss
model = EmbeddingUnet(setup_config)
loss = EmbeddingLoss(setup_config)
if setup_config.get("RADAM"):
optimizer = RAdam(model.parameters(),
lr=0.5e-5,
betas=(0.95, 0.999),
eps=1e-8)
else:
optimizer = torch.optim.Adam(model.parameters(),
lr=0.5e-5,
betas=(0.95, 0.999),
eps=1e-8)
if setup_config["AUX_TASK"]:
pipeline = (pipeline + ToInt64(gt_labels) + UnSqueeze(
[raw, gt_labels, mask, neighborhood_gt, neighborhood_mask]) +
TrainEmbedding(
model=model,
optimizer=optimizer,
loss=loss,
inputs={"raw": raw},
loss_inputs={
0: embedding,
"target": gt_labels,
"mask": mask,
"neighborhood": neighborhood_gt,
"neighborhood_mask": neighborhood_mask,
"neighborhood_target": neighborhood,
},
outputs={
0: embedding,
1: neighborhood
},
gradients={
0: embedding_gradient,
1: neighborhood_gradient
},
array_specs={
embedding:
ArraySpec(dtype=np.float32, voxel_size=voxel_size),
embedding_gradient:
ArraySpec(dtype=np.float32, voxel_size=voxel_size),
neighborhood:
ArraySpec(dtype=np.float32, voxel_size=voxel_size),
neighborhood_gradient:
ArraySpec(dtype=np.float32, voxel_size=voxel_size),
},
save_every=checkpoint_every,
log_dir=tensorboard_log_dir,
checkpoint_basename=embedding_net_name,
) + Squeeze([
embedding,
embedding_gradient,
neighborhood,
neighborhood_gradient,
neighborhood_gt,
neighborhood_mask,
raw,
gt_labels,
mask,
]))
return (
pipeline,
embedding,
embedding_gradient,
neighborhood,
neighborhood_gradient,
)
else:
pipeline = (
pipeline + ToInt64(gt_labels) + UnSqueeze([raw, gt_labels, mask]) +
TrainEmbedding(
model=model,
optimizer=optimizer,
loss=loss,
inputs={"raw": raw},
loss_inputs={
0: embedding,
"target": gt_labels,
"mask": mask
},
outputs={0: embedding},
gradients={0: embedding_gradient},
array_specs={
embedding:
ArraySpec(dtype=np.float32, voxel_size=voxel_size),
embedding_gradient:
ArraySpec(dtype=np.float32, voxel_size=voxel_size),
},
save_every=checkpoint_every,
log_dir=tensorboard_log_dir,
checkpoint_basename=embedding_net_name,
) + Squeeze([embedding, embedding_gradient, raw, gt_labels, mask]))
return pipeline, embedding, embedding_gradient
def __read_spec(self, array_key, base_folder, ds_name):
js_path = os.path.join(base_folder, "%s.json" % ds_name)
if not os.path.exists(js_path):
raise RuntimeError("%s.json not in %s" % (ds_name, base_folder))
dataset_spec = json.load(open(js_path, 'r'))
self.sec_dir = dataset_spec['sections_dir']
self.sections = dataset_spec['sections']
dims = np.asarray(dataset_spec['shape'])
self.tile_shape = np.asarray(dataset_spec['tile_size'])
# last two dims are assumed to be rows x columns
dims[-2:] *= self.tile_shape
if self.ndims is None:
self.ndims = len(dims)
else:
assert self.ndims == len(dims)
if array_key in self.array_specs:
spec = self.array_specs[array_key].copy()
else:
spec = ArraySpec()
if spec.voxel_size is None:
if 'resolution' in dataset_spec:
spec.voxel_size = Coordinate(dataset_spec['resolution'])
else:
spec.voxel_size = Coordinate((1, ) * self.ndims)
logger.warning(
"WARNING: File %s does not contain resolution information "
"for %s (dataset %s), voxel size has been set to %s. This "
"might not be what you want.", self.filename, array_key,
ds_name, spec.voxel_size)
if spec.roi is None:
if 'offset' in dataset_spec:
offset = Coordinate(dataset_spec['offset'])
else:
offset = Coordinate((0, ) * self.ndims)
spec.roi = Roi(offset, dims * spec.voxel_size)
if spec.dtype is not None:
assert spec.dtype == dataset.dtype, (
"dtype %s provided in array_specs for %s, "
"but differs from dataset %s dtype %s" %
(self.array_specs[array_key].dtype, array_key, ds_name,
dataset_spec['dtype']))
else:
spec.dtype = dataset_spec['dtype']
if spec.interpolatable is None:
spec.interpolatable = spec.dtype in [
np.float,
np.float32,
np.float64,
np.float128,
np.uint8 # assuming this is not used for labels
]
logger.warning(
"WARNING: You didn't set 'interpolatable' for %s "
"(dataset %s). Based on the dtype %s, it has been "
"set to %s. This might not be what you want.", array_key,
ds_name, spec.dtype, spec.interpolatable)
return spec
