def _query_kdtree(self, node: cKDTreeNode,
bb: Tuple[np.ndarray, np.ndarray]) -> List[np.ndarray]:
def substitute_dim(bound: np.ndarray, sub_dim: int, sub: float):
# replace bound[sub_dim] with sub
return np.array(
[bound[i] if i != sub_dim else sub for i in range(3)])
if node is None:
return []
if node.split_dim != -1:
# if the split location is above the roi, no need to split (cannot be above None)
if bb[1][node.split_dim] is not None and node.split > bb[1][
node.split_dim]:
return self._query_kdtree(node.lesser, bb)
elif (bb[0][node.split_dim] is not None
and node.split < bb[0][node.split_dim]):
return self._query_kdtree(node.greater, bb)
else:
return self._query_kdtree(
node.greater, bb) + self._query_kdtree(node.lesser, bb)
else:
# handle leaf node
bbox = Roi(
Coordinate(bb[0]),
Coordinate(
tuple(y - x if x is not None and y is not None else y
for x, y in zip(*bb))),
)
points = [
ind for ind, point in zip(node.indices, node.data_points)
if bbox.contains(np.round(point))
]
return points
def _resample_relative(
self, inside: np.ndarray, outside: np.ndarray, bb: Roi
) -> Optional[np.ndarray]:
offset = outside - inside
with np.errstate(divide="ignore", invalid="ignore"):
# bb_crossings will be 0 if inside is on the bb, 1 if outside is on the bb
bb_x = np.asarray(
[
(np.asarray(bb.get_begin()) - inside) / offset,
(np.asarray(bb.get_end() - Coordinate([1, 1, 1])) - inside)
/ offset,
]
)
if np.sum(np.logical_and((bb_x > 0), (bb_x <= 1))) > 0:
# all values of bb_x between 0, 1 represent a crossing of a bounding plane
# the minimum of which is the (normalized) distance to the closest bounding plane
s = np.min(bb_x[np.logical_and((bb_x > 0), (bb_x <= 1))])
return Coordinate(np.floor(np.array(inside) + s * offset))
else:
logging.debug(
(
"Could not create a node on the bounding box {} "
+ "given points (inside:{}, ouside:{})"
).format(bb, inside, outside)
)
return None
def __get_spec(self, array_key):
info = self.__get_info(array_key)
roi_min = info['Extended']['MinPoint']
if roi_min is not None:
roi_min = Coordinate(roi_min[::-1])
roi_max = info['Extended']['MaxPoint']
if roi_max is not None:
roi_max = Coordinate(roi_max[::-1])
data_roi = Roi(offset=roi_min, shape=(roi_max - roi_min))
data_dims = Coordinate(data_roi.get_shape())
if self.ndims is None:
self.ndims = len(data_dims)
else:
assert self.ndims == len(data_dims)
if array_key in self.array_specs:
spec = self.array_specs[array_key].copy()
else:
spec = ArraySpec()
if spec.voxel_size is None:
spec.voxel_size = Coordinate(info['Extended']['VoxelSize'])
if spec.roi is None:
spec.roi = data_roi * spec.voxel_size
data_dtype = dvision.DVIDDataInstance(self.hostname, self.port,
self.uuid,
self.datasets[array_key]).dtype
if spec.dtype is not None:
assert spec.dtype == data_dtype, (
"dtype %s provided in array_specs for %s, "
"but differs from instance %s dtype %s" %
(self.array_specs[array_key].dtype, array_key,
self.datasets[array_key], data_dtype))
else:
spec.dtype = data_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.", array_key, spec.dtype,
spec.interpolatable)
return spec
def process(self, batch, request):
points = batch.points[self.points]
voxel_size = self.spec[self.array].voxel_size
# get roi used for creating the new array (points_roi does not
# necessarily align with voxel size)
enlarged_vol_roi = points.spec.roi.snap_to_grid(voxel_size)
offset = enlarged_vol_roi.get_begin() / voxel_size
shape = enlarged_vol_roi.get_shape() / voxel_size
data_roi = Roi(offset, shape)
# points ROI is at least +- 1 in t of requested array ROI, we can save
# some time by shaving the excess off
data_roi = data_roi.grow((-1, 0, 0, 0), (-1, 0, 0, 0))
logger.debug("Points in %s", points.spec.roi)
for i, point in points.data.items():
logger.debug("%d, %s", i, point.location)
logger.debug("Data roi in voxels: %s", data_roi)
logger.debug("Data roi in world units: %s", data_roi * voxel_size)
parent_vectors_data, mask_data = self.__draw_parent_vectors(
points, data_roi, voxel_size, enlarged_vol_roi.get_begin(),
self.radius)
# create array and crop it to requested roi
spec = self.spec[self.array].copy()
spec.roi = data_roi * voxel_size
parent_vectors = Array(data=parent_vectors_data, spec=spec)
logger.debug("Cropping parent vectors to %s", request[self.array].roi)
batch.arrays[self.array] = parent_vectors.crop(request[self.array].roi)
# create mask and crop it to requested roi
spec = self.spec[self.mask].copy()
spec.roi = data_roi * voxel_size
mask = Array(data=mask_data, spec=spec)
logger.debug("Cropping mask to %s", request[self.mask].roi)
batch.arrays[self.mask] = mask.crop(request[self.mask].roi)
# restore requested ROI of points
if self.points in request:
request_roi = request[self.points].roi
points.spec.roi = request_roi
for i, p in list(points.data.items()):
if not request_roi.contains(p.location):
del points.data[i]
if len(points.data) == 0:
logger.warning("Returning empty batch for key %s and roi %s" %
(self.points, request_roi))
def prepare(self, request):
logger.debug("request: %s"%request)
logger.debug("upstream spec: %s"%self.upstream_spec)
# remember request
self.request = copy.deepcopy(request)
for volume_type in self.pad_sizes.keys():
if volume_type not in request.volumes:
continue
roi = request.volumes[volume_type]
# check out-of-bounds
# TODO: this should be moved to super class, this should hold for any
# batch provider
if self.pad_sizes[volume_type] is not None:
if not self.spec.volumes[volume_type].intersects(roi):
raise RuntimeError("%s ROI %s lies outside of padded ROI %s"%(volume_type,roi,self.spec.volumes[volume_type]))
# change request to fit into upstream spec
request.volumes[volume_type] = roi.intersect(self.upstream_spec.volumes[volume_type])
if request.volumes[volume_type] is None:
logger.warning("Requested %s ROI lies entirely outside of upstream ROI."%volume_type)
# ensure a valid request by asking for empty ROI
request.volumes[volume_type] = Roi(
self.upstream_spec.volumes[volume_type].get_offset(),
(0,)*self.upstream_spec.volumes[volume_type].dims()
)
logger.debug("new request: %s"%request)
def setup(self):
f = h5py.File(self.filename, 'r')
self.spec = ProviderSpec()
self.ndims = None
for (volume_type, ds) in self.datasets.items():
if ds not in f:
raise RuntimeError("%s not in %s" % (ds, self.filename))
dims = f[ds].shape
self.spec.volumes[volume_type] = Roi((0, ) * len(dims), dims)
if self.ndims is None:
self.ndims = len(dims)
else:
assert self.ndims == len(dims)
if self.specified_resolution is None:
if 'resolution' in f[ds].attrs:
self.resolutions[volume_type] = tuple(
f[ds].attrs['resolution'])
else:
default_resolution = (1, ) * self.ndims
logger.warning(
"WARNING: your source does not contain resolution information"
" (no attribute 'resolution' in {} dataset). I will assume {}. "
"This might not be what you want.".format(
ds, default_resolution))
self.resolutions[volume_type] = default_resolution
else:
self.resolutions[volume_type] = self.specified_resolution
f.close()
def prepare(self, request):
upstream_spec = self.get_upstream_provider().spec
logger.debug("request: %s" % request)
logger.debug("upstream spec: %s" % upstream_spec)
# TODO: remove this?
if self.key not in request:
return
roi = request[self.key].roi.copy()
# change request to fit into upstream spec
request[self.key].roi = roi.intersect(upstream_spec[self.key].roi)
if request[self.key].roi.empty():
logger.warning(
"Requested %s ROI %s lies entirely outside of upstream "
"ROI %s.", self.key, roi, upstream_spec[self.key].roi)
# ensure a valid request by asking for empty ROI
request[self.key].roi = Roi(
upstream_spec[self.key].roi.get_offset(),
(0, ) * upstream_spec[self.key].roi.dims())
logger.debug("new request: %s" % request)
deps = BatchRequest()
deps[self.key] = request[self.key]
return deps
def __read_spec(self, cv):
# dataset = data_file[ds_name]
dims = Coordinate(cv.bounds.maxpt[::-1]) * self._get_voxel_size(cv)
if self.ndims is None:
self.ndims = cv.bounds.ndim
else:
assert self.ndims == len(dims)
if self.array_spec is not None:
spec = self.array_spec.copy()
else:
spec = ArraySpec()
if spec.voxel_size is None:
voxel_size = self._get_voxel_size(cv)
if voxel_size is None:
voxel_size = Coordinate((1,) * self.ndims)
logger.warning(
"WARNING: File %s does not contain resolution information "
"for %s , voxel size has been set to %s. This "
"might not be what you want.",
self.filename, array_key, spec.voxel_size)
spec.voxel_size = voxel_size
if spec.roi is None:
offset = self._get_offset(cv)
if offset is None:
offset = Coordinate((0,) * self.ndims)
spec.roi = Roi(offset, dims * spec.voxel_size)
if spec.dtype is not None:
assert spec.dtype == cv.dtype, (
"dtype %s provided in array_specs for %s, "
"but differs from cloudvolume dtype %s" %
(self.array_spec.dtype,
self.array_key, dataset.dtype))
else:
spec.dtype = cv.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_key, spec.dtype,
spec.interpolatable)
return spec
def __read_spec(self, array_key, data_file, ds_name):
dataset = data_file[ds_name]
if array_key in self.array_specs:
spec = self.array_specs[array_key].copy()
else:
spec = ArraySpec()
if spec.voxel_size is None:
voxel_size = self._get_voxel_size(dataset)
if voxel_size is None:
voxel_size = Coordinate((1, ) * len(dataset.shape))
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)
spec.voxel_size = voxel_size
self.ndims = len(spec.voxel_size)
if spec.roi is None:
offset = self._get_offset(dataset)
if offset is None:
offset = Coordinate((0, ) * self.ndims)
if self.channels_first:
shape = Coordinate(dataset.shape[-self.ndims:])
else:
shape = Coordinate(dataset.shape[:self.ndims])
spec.roi = Roi(offset, shape * 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 __get_chunks(self):
daisy_client = daisy.Client()
while True:
block = daisy_client.acquire_block()
if block is None:
return
logger.info("Processing block %s", block)
start = time.time()
chunk_request = self.reference.copy()
for key, reference_spec in self.reference.items():
roi_type = self.roi_map.get(key, None)
if roi_type is None:
raise RuntimeError(
"roi_map does not map item %s to either 'read_roi' "
"or 'write_roi'" % key)
if roi_type == 'read_roi':
chunk_request[key].roi = Roi(
block.read_roi.get_offset(),
block.read_roi.get_shape())
elif roi_type == 'write_roi':
chunk_request[key].roi = Roi(
block.write_roi.get_offset(),
block.write_roi.get_shape())
else:
raise RuntimeError(
"%s is not a vaid ROI type (read_roi or write_roi)")
self.get_upstream_provider().request_batch(chunk_request)
end = time.time()
if self.block_done_callback:
self.block_done_callback(block, start, end - start)
daisy_client.release_block(block, ret=0)
def __get_roi(self, array_name):
data_instance = dvision.DVIDDataInstance(self.hostname, self.port, self.uuid, array_name)
info = data_instance.info
roi_min = info['Extended']['MinPoint']
if roi_min is not None:
roi_min = Coordinate(roi_min[::-1])
roi_max = info['Extended']['MaxPoint']
if roi_max is not None:
roi_max = Coordinate(roi_max[::-1])
return Roi(offset=roi_min, shape=roi_max - roi_min)
def __init__(
self,
dbname: str,
url: str,
points: List[GraphKey],
graph_specs: Optional[Union[GraphSpec, List[GraphSpec]]] = None,
directed: bool = False,
total_roi: Roi = None,
nodes_collection: str = "nodes",
edges_collection: str = "edges",
meta_collection: str = "meta",
endpoint_names: Tuple[str, str] = ("u", "v"),
position_attribute: str = "position",
node_attrs: Optional[List[str]] = None,
edge_attrs: Optional[List[str]] = None,
nodes_filter: Optional[Dict[str, Any]] = None,
edges_filter: Optional[Dict[str, Any]] = None,
edge_inclusion: str = "either",
node_inclusion: str = "dangling",
fail_on_inconsistent_node: bool = False,
):
self.points = points
graph_specs = (graph_specs if graph_specs is not None else GraphSpec(
Roi(Coordinate([None] * 3), Coordinate([None] * 3)),
directed=False))
specs = (graph_specs if isinstance(graph_specs, list)
and len(graph_specs) == len(points) else [graph_specs] *
len(points))
self.specs = {key: spec for key, spec in zip(points, specs)}
self.directed = directed
self.nodes_collection = nodes_collection
self.edges_collection = edges_collection
self.meta_collection = meta_collection
self.endpoint_names = endpoint_names
self.position_attribute = position_attribute
self.position_attribute = position_attribute
self.node_attrs = node_attrs
self.edge_attrs = edge_attrs
self.nodes_filter = nodes_filter
self.edges_filter = edges_filter
self.edge_inclusion = edge_inclusion
self.node_inclusion = node_inclusion
self.dbname = dbname
self.url = url
self.nodes_collection = nodes_collection
self.fail_on_inconsistent_node = fail_on_inconsistent_node
self.graph_provider = None
def process(self, batch, request):
src_points = batch.points[self.src_points]
voxel_size = self.spec[self.array].voxel_size
# get roi used for creating the new array (points_roi does not
# necessarily align with voxel size)
enlarged_vol_roi = src_points.spec.roi.snap_to_grid(voxel_size)
offset = enlarged_vol_roi.get_begin() / voxel_size
shape = enlarged_vol_roi.get_shape() / voxel_size
data_roi = Roi(offset, shape)
logger.debug("Src points in %s", src_points.spec.roi)
for i, point in src_points.data.items():
logger.debug("%d, %s", i, point.location)
logger.debug("Data roi in voxels: %s", data_roi)
logger.debug("Data roi in world units: %s", data_roi * voxel_size)
mask_array = None if self.mask is None else batch.arrays[
self.mask].crop(enlarged_vol_roi)
partner_vectors_data, pointmask = self.__draw_partner_vectors(
src_points, batch.points[self.trg_points], data_roi, voxel_size,
enlarged_vol_roi.get_begin(), self.radius, mask_array)
# create array and crop it to requested roi
spec = self.spec[self.array].copy()
spec.roi = data_roi * voxel_size
partner_vectors = Array(data=partner_vectors_data, spec=spec)
logger.debug("Cropping partner vectors to %s", request[self.array].roi)
batch.arrays[self.array] = partner_vectors.crop(
request[self.array].roi)
if self.pointmask is not None and self.pointmask in request:
spec = self.spec[self.array].copy()
spec.roi = data_roi * voxel_size
pointmask = Array(data=np.array(pointmask, dtype=spec.dtype),
spec=spec)
batch.arrays[self.pointmask] = pointmask.crop(
request[self.pointmask].roi)
# restore requested ROI of src and target points.
if self.src_points in request:
self.__restore_points_roi(request, self.src_points,
batch.points[self.src_points])
if self.trg_points in request:
self.__restore_points_roi(request, self.trg_points,
batch.points[self.trg_points])
# restore requested objectmask
if self.mask is not None:
batch.arrays[self.mask] = batch.arrays[self.mask].crop(
request[self.mask].roi)
def __init__(
self,
dbname: str,
url: str,
points: List[GraphKey],
graph_specs: Optional[Union[GraphSpec, List[GraphSpec]]] = None,
directed: bool = False,
total_roi: Roi = None,
nodes_collection: str = "nodes",
edges_collection: str = "edges",
meta_collection: str = "meta",
endpoint_names: Tuple[str, str] = ("u", "v"),
position_attribute: str = "position",
node_attrs: Optional[List[str]] = None,
edge_attrs: Optional[List[str]] = None,
nodes_filter: Optional[Dict[str, Any]] = None,
edges_filter: Optional[Dict[str, Any]] = None,
num_nodes=100000,
dist_attribute=None,
min_dist=29000,
):
self.points = points
graph_specs = (graph_specs if graph_specs is not None else GraphSpec(
Roi(Coordinate([None] * 3), Coordinate([None] * 3)),
directed=False))
specs = (graph_specs if isinstance(graph_specs, list)
and len(graph_specs) == len(points) else [graph_specs] *
len(points))
self.specs = {key: spec for key, spec in zip(points, specs)}
self.position_attribute = position_attribute
self.node_attrs = node_attrs
self.edge_attrs = edge_attrs
self.nodes_filter = nodes_filter
self.edges_filter = edges_filter
self.dist_attribute = dist_attribute
self.min_dist = min_dist
self.num_nodes = num_nodes
self.graph_provider = MongoDbGraphProvider(
dbname,
url,
mode="r+",
directed=directed,
total_roi=None,
nodes_collection=nodes_collection,
edges_collection=edges_collection,
meta_collection=meta_collection,
endpoint_names=endpoint_names,
position_attribute=position_attribute,
)
def _query_kdtree(
self, node: cKDTreeNode, bb: Tuple[np.ndarray, np.ndarray]
) -> List[np.ndarray]:
def substitute_dim(bound: np.ndarray, sub_dim: int, sub: float):
# replace bound[sub_dim] with sub
return np.array([bound[i] if i != sub_dim else sub for i in range(3)])
if node is None:
return []
if node.split_dim != -1:
# recursive handling of child nodes
greater_roi = (substitute_dim(bb[0], node.split_dim, node.split), bb[1])
lesser_roi = (bb[0], substitute_dim(bb[1], node.split_dim, node.split))
return self._query_kdtree(node.greater, greater_roi) + self._query_kdtree(
node.lesser, lesser_roi
)
else:
# handle leaf node
# TODO: handle bounding box properly. bb[0], and bb[1] may not be integers.
bbox = Roi(Coordinate(bb[0]), Coordinate(bb[1] - bb[0]))
points = [point for point in node.data_points if bbox.contains(point)]
return points
def __get_source_roi(self, transformation):
dims = transformation.shape[0]
# get bounding box of needed data for transformation
bb_min = Coordinate(
int(math.floor(transformation[d].min())) for d in range(dims))
bb_max = Coordinate(
int(math.ceil(transformation[d].max())) + 1 for d in range(dims))
# create roi sufficiently large to feed transformation
source_roi = Roi(bb_min, bb_max - bb_min)
return source_roi
def compute_upstream_roi(request_roi, sub_shift_array):
""" Compute the ROI to pass upstream for a specific item (array or points) in a request
:param request_roi: the downstream ROI passed to the Jitter node
:param sub_shift_array: the portion of the global shift array that should be used to shift the item
:return: the expanded ROI to pass upstream
"""
max_shift = Coordinate(sub_shift_array.max(axis=0))
min_shift = Coordinate(sub_shift_array.min(axis=0))
downstream_offset = request_roi.get_offset()
upstream_offset = downstream_offset - max_shift
upstream_shape = request_roi.get_shape() + max_shift - min_shift
return Roi(offset=upstream_offset, shape=upstream_shape)
def setup(self):
self._read_points()
if self.points_spec is not None:
self.provides(self.points, self.points_spec)
return
min_bb = Coordinate(np.floor(np.amin(self.data[:, :self.ndims], 0)))
max_bb = Coordinate(np.ceil(np.amax(self.data[:, :self.ndims], 0)) + 1)
roi = Roi(min_bb, max_bb - min_bb)
self.provides(self.points, GraphSpec(roi=roi))
def setup(self):
self.data = self.read_points(self.filename)
self.ndims = self.data.shape[1]
if self.points_spec is not None:
self.provides(self.points, self.points_spec)
return
min_bb = Coordinate(np.floor(np.amin(self.data, 0)))
max_bb = Coordinate(np.ceil(np.amax(self.data, 0)))
roi = Roi(min_bb, max_bb - min_bb)
self.provides(self.points, PointsSpec(roi=roi))
def __read_spec(self, array_key, data_file, ds_name):
dataset = data_file[ds_name]
if array_key in self.array_specs:
spec = self.array_specs[array_key].copy()
else:
spec = ArraySpec()
if spec.voxel_size is None:
voxel_size = self._get_voxel_size(dataset)
if voxel_size is None:
voxel_size = Coordinate((1, ) * len(dataset.shape))
spec.voxel_size = voxel_size
self.ndims = len(spec.voxel_size)
if spec.roi is None:
offset = self._get_offset(dataset)
if offset is None:
offset = Coordinate((0, ) * self.ndims)
if self.data_format == 'channels_first':
shape = Coordinate(dataset.shape[-self.ndims:])
if self.data_format == 'channels_last':
shape = Coordinate(dataset.shape[0:self.ndims])
spec.roi = Roi(offset, shape * 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
]
return spec
def process(self, batch, request):
filename = os.path.join(self.output_dir, self.output_filename)
if not self.dataset_offsets:
self.init_datasets(batch)
with self._open_file(filename) as data_file:
for (array_key, dataset_name) in self.dataset_names.items():
dataset = data_file[dataset_name]
array_roi = batch.arrays[array_key].spec.roi
voxel_size = self.spec[array_key].voxel_size
dims = array_roi.dims()
channel_slices = (slice(None), ) * max(
0,
len(dataset.shape) - dims)
dataset_roi = Roi(
self.dataset_offsets[array_key],
Coordinate(dataset.shape[-dims:]) * voxel_size)
common_roi = array_roi.intersect(dataset_roi)
if common_roi.empty():
logger.warn(
"array %s with ROI %s lies outside of dataset ROI %s, "
"skipping writing" %
(array_key, array_roi, dataset_roi))
continue
dataset_voxel_roi = (
common_roi - self.dataset_offsets[array_key]) // voxel_size
dataset_voxel_slices = dataset_voxel_roi.to_slices()
array_voxel_roi = (common_roi -
array_roi.get_offset()) // voxel_size
array_voxel_slices = array_voxel_roi.to_slices()
logger.debug("writing %s to voxel coordinates %s" %
(array_key, dataset_voxel_roi))
data = batch.arrays[array_key].data[channel_slices +
array_voxel_slices]
dataset[channel_slices + dataset_voxel_slices] = data
def __predict(self, use_gpu):
if not self.net_initialized:
logger.info("Initializing solver...")
if use_gpu is not None:
logger.debug("Predict process: using GPU %d" % use_gpu)
caffe.enumerate_devices(False)
caffe.set_devices((use_gpu, ))
caffe.set_mode_gpu()
caffe.select_device(use_gpu, False)
self.net = caffe.Net(self.prototxt, self.weights, caffe.TEST)
self.net_io = NetIoWrapper(self.net)
self.net_initialized = True
start = time.time()
batch = self.batch_in.get()
fetch_time = time.time() - start
self.net_io.set_inputs({
'data':
batch.volumes[VolumeTypes.RAW].data[np.newaxis, np.newaxis, :],
})
self.net.forward()
output = self.net_io.get_outputs()
predict_time = time.time() - start
logger.info(
"Predict process: time=%f (including %f waiting for batch)" %
(predict_time, fetch_time))
assert len(
output['aff_pred'].shape
) == 5, "Got affinity prediction with unexpected number of dimensions, should be 1 (direction) + 3 (spatial) + 1 (batch, not used), but is %d" % len(
output['aff_pred'].shape)
batch.volumes[VolumeTypes.PRED_AFFINITIES] = Volume(
output['aff_pred'][0], Roi(), (1, 1, 1))
return batch
def setup(self):
logger.debug("Initializing graph source")
self._read_points()
logger.debug("Graph source initialized")
if self.points_spec is not None:
assert len(self.points) == len(self.points_spec)
for point, point_spec in zip(self.points, self.points_spec):
self.provides(point, point_spec)
else:
logger.debug("No point spec provided!")
roi = Roi(Coordinate((None, None, None)),
Coordinate((None, None, None)))
for point in self.points:
self.provides(point, PointsSpec(roi=roi))
def __recompute_roi(self, roi, transformation):
dims = roi.dims()
# get bounding box of needed data for transformation
bb_min = Coordinate(
int(math.floor(transformation[d].min())) for d in range(dims))
bb_max = Coordinate(
int(math.ceil(transformation[d].max())) + 1 for d in range(dims))
# create roi sufficiently large to feed transformation
source_roi = Roi(bb_min, bb_max - bb_min)
# shift transformation, such that it can be applied on indices of source
# batch
for d in range(dims):
transformation[d] -= bb_min[d]
return source_roi
def setup(self):
self._read_points()
if self.points_spec is not None:
self.provides(self.points, self.points_spec)
else:
# If you don't provide a queryset, the roi will shrink to fit the
# points in the swc(s). This may cause problems if you expect empty point
# sets when querying the edges of your volume
logging.warning("No point spec provided!")
min_bb = Coordinate(self.data.mins[0:3])
# cKDTree max is inclusive
max_bb = Coordinate(self.data.maxes[0:3]) + Coordinate([1, 1, 1])
roi = Roi(min_bb, max_bb - min_bb)
self.provides(self.points, PointsSpec(roi=roi))
def __read(self, roi):
if len(self.files) == 1:
return self.__read_file(self.files[0], roi)
else:
file_indices = range(
roi.get_begin()[0],
roi.get_end()[0])
file_roi = Roi(
roi.get_begin()[1:],
roi.get_shape()[1:])
return np.array([
self.__read_file(self.files[i], file_roi)
for i in file_indices
])
def provide(self, request: BatchRequest) -> Batch:
random.seed(request.random_seed)
np.random.seed(request.random_seed)
timing = Timing(self, "provide")
timing.start()
batch = Batch()
roi = request[self.points].roi
region_shape = roi.get_shape()
trees = []
for _ in range(self.n_obj):
for _ in range(100):
root = np.random.random(len(region_shape)) * region_shape
tree = self._grow_tree(
root, Roi((0,) * len(region_shape), region_shape)
)
if self.num_nodes[0] <= len(tree.nodes) <= self.num_nodes[1]:
break
trees.append(tree)
# logger.info("{} trees got, expected {}".format(len(trees), self.n_obj))
trees_graph = nx.disjoint_union_all(trees)
points = {
node_id: Node(np.floor(node_attrs["pos"]) + roi.get_begin())
for node_id, node_attrs in trees_graph.nodes.items()
}
batch[self.points] = Graph(points, request[self.points], list(trees_graph.edges))
timing.stop()
batch.profiling_stats.add(timing)
# self._plot_tree(tree)
return batch
def setup(self):
hdf_file = h5py.File(self.filename, 'r')
for (volume_type, ds_name) in self.datasets.items():
if ds_name not in hdf_file:
raise RuntimeError("%s not in %s" % (ds_name, self.filename))
spec = self.__read_spec(volume_type, hdf_file, ds_name)
self.provides(volume_type, spec)
if self.points_types is not None:
for points_type in self.points_types:
spec = PointsSpec()
spec.roi = Roi(self.points_rois[points_type][0],
self.points_rois[points_type][1])
self.provides(points_type, spec)
hdf_file.close()
def _grow_tree(self, root: np.ndarray, roi: Roi) -> nx.DiGraph:
tree = nx.DiGraph()
tree.add_node(0, pos=root)
leaves = deque([0])
next_id = 1
while len(leaves) > 0 and len(tree.nodes) < self.num_nodes[1]:
current_id = leaves.pop()
current_loc = tree.nodes[current_id]["pos"]
tail = [current_loc]
while len(tail) < self.ma and len(list(tree.predecessors(current_id))) == 1:
preds = list(tree.predecessors(current_id))
tail.append(tree.nodes[preds[0]]["pos"])
if len(tail) == 1:
tail.append(np.random.random(tail[0].shape)-0.5 + tail[0])
moving_direction = np.array(tail[0] - tail[-1])
if len(tree.nodes) >= self.num_nodes[0] and random.random() < self.p_die:
continue
else:
num_branches = np.random.choice(
range(1, self.max_split + 1), p=self.split_ps
)
next_points = self._gen_n(num_branches, current_loc, moving_direction)
for point in next_points:
if not roi.contains(point):
continue
tree.add_node(next_id, pos=point)
tree.add_edge(current_id, next_id)
leaves.appendleft(next_id)
next_id += 1
return tree
def __select_random_location_with_points(
self,
request,
lcm_shift_roi,
lcm_voxel_size):
request_points_roi = request[self.ensure_nonempty].roi
while True:
# How to pick shifts that ensure that a randomly chosen point is
# contained in the request ROI:
#
#
# request point
# [---------) .
# 0 +10 17
#
# least shifted to contain point
# [---------)
# 8 +10
# ==
# point-request.begin-request.shape+1
#
# most shifted to contain point:
# [---------)
# 17 +10
# ==
# point-request.begin
#
# all possible shifts
# [---------)
# 8 +10
# ==
# point-request.begin-request.shape+1
# ==
# request.shape
#
# In the most shifted scenario, it could happen that the point lies
# exactly at the lower boundary (17 in the example). This will cause
# trouble if later we mirror the batch. The point would end up lying
# on the other boundary, which is exclusive and thus not part of the
# ROI. Therefore, we have to ensure that the point is well inside
# the shifted ROI, not just on the boundary:
#
# all possible shifts
# [--------)
# 8 +9
# ==
# request.shape-1
# pick a random point
point_id = choice(self.points.data.keys())
point = self.points.data[point_id]
logger.debug(
"select random point %d at %s",
point_id,
point.location)
# get the lcm voxel that contains this point
lcm_location = Coordinate(point.location/lcm_voxel_size)
logger.debug(
"belongs to lcm voxel %s",
lcm_location)
# mark all dimensions in which the point lies on the lower boundary
# of the lcm voxel
on_lower_boundary = lcm_location*lcm_voxel_size == point.location
logger.debug(
"lies on the lower boundary of the lcm voxel in dimensions %s",
on_lower_boundary)
# for each of these dimensions, we have to change the shape of the
# shift ROI using the following correction
lower_boundary_correction = Coordinate((
-1 if o else 0
for o in on_lower_boundary
))
logger.debug(
"lower bound correction for shape of shift ROI %s",
lower_boundary_correction)
# get the request ROI's shape in lcm
lcm_roi_begin = request_points_roi.get_begin()/lcm_voxel_size
lcm_roi_shape = request_points_roi.get_shape()/lcm_voxel_size
logger.debug("Point request ROI: %s", request_points_roi)
logger.debug("Point request lcm ROI shape: %s", lcm_roi_shape)
# get all possible starting points of lcm_roi_shape that contain
# lcm_location
lcm_shift_roi_begin = (
lcm_location - lcm_roi_begin - lcm_roi_shape +
Coordinate((1,)*len(lcm_location))
)
lcm_shift_roi_shape = (
lcm_roi_shape + lower_boundary_correction
)
lcm_point_shift_roi = Roi(lcm_shift_roi_begin, lcm_shift_roi_shape)
logger.debug("lcm point shift roi: %s", lcm_point_shift_roi)
# intersect with total shift ROI
if not lcm_point_shift_roi.intersects(lcm_shift_roi):
logger.debug(
"reject random shift, random point %s shift ROI %s does "
"not intersect total shift ROI %s", point.location,
lcm_point_shift_roi, lcm_shift_roi)
continue
lcm_point_shift_roi = lcm_point_shift_roi.intersect(lcm_shift_roi)
# select a random shift from all possible shifts
random_shift = self.__select_random_location(
lcm_point_shift_roi,
lcm_voxel_size)
logger.debug("random shift: %s", random_shift)
# count all points inside the shifted ROI
points_request = BatchRequest()
points_request[self.ensure_nonempty] = PointsSpec(
roi=request_points_roi.shift(random_shift))
logger.debug("points request: %s", points_request)
points_batch = self.get_upstream_provider().request_batch(points_request)
point_ids = points_batch.points[self.ensure_nonempty].data.keys()
assert point_id in point_ids, (
"Requested batch to contain point %s, but got points "
"%s"%(point_id, point_ids))
num_points = len(point_ids)
# accept this shift with p=1/num_points
#
# This is to compensate the bias introduced by close-by points.
accept = random() <= 1.0/num_points
if accept:
return random_shift