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