def test_output(self): a = ArrayKey("A") b = ArrayKey("B") source_a = TestSourceRandomLocation(a) source_b = TestSourceRandomLocation(b) pipeline = (source_a, source_b) + \ MergeProvider() + CustomRandomLocation() with build(pipeline): for i in range(10): batch = pipeline.request_batch( BatchRequest({ a: ArraySpec(roi=Roi((0, 0, 0), (20, 20, 20))), b: ArraySpec(roi=Roi((0, 0, 0), (20, 20, 20))) })) self.assertTrue(np.sum(batch.arrays[a].data) > 0) self.assertTrue(np.sum(batch.arrays[b].data) > 0) # Request a ROI with the same shape as the entire ROI full_roi_a = Roi((0, 0, 0), source_a.roi.get_shape()) full_roi_b = Roi((0, 0, 0), source_b.roi.get_shape()) batch = pipeline.request_batch( BatchRequest({ a: ArraySpec(roi=full_roi_a), b: ArraySpec(roi=full_roi_b) }))
def test_placeholder(self): test_labels = ArrayKey("TEST_LABELS") test_points = GraphKey("TEST_POINTS") pipeline = ( PointTestSource3D() + RandomLocation(ensure_nonempty=test_points) + ElasticAugment([10, 10, 10], [0.1, 0.1, 0.1], [0, 2.0 * math.pi]) + Snapshot( {test_labels: "volumes/labels"}, output_dir=self.path_to(), output_filename="elastic_augment_test{id}-{iteration}.hdf", )) with build(pipeline): for i in range(2): request_size = Coordinate((40, 40, 40)) request_a = BatchRequest(random_seed=i) request_a.add(test_points, request_size) request_a.add(test_labels, request_size, placeholder=True) request_b = BatchRequest(random_seed=i) request_b.add(test_points, request_size) request_b.add(test_labels, request_size) batch_a = pipeline.request_batch(request_a) batch_b = pipeline.request_batch(request_b) points_a = batch_a[test_points].nodes points_b = batch_b[test_points].nodes for a, b in zip(points_a, points_b): assert all(np.isclose(a.location, b.location))
def test_ensure_center_non_zero(self): path = Path(self.path_to("test_swc_source.swc")) # write test swc self._write_swc(path, self._toy_swc_points().to_nx_graph()) # read arrays swc = PointsKey("SWC") img = ArrayKey("IMG") pipeline = (SwcFileSource( path, [swc], [PointsSpec(roi=Roi((0, 0, 0), (11, 11, 11)))]) + RandomLocation(ensure_nonempty=swc, ensure_centered=True) + RasterizeSkeleton( points=swc, array=img, array_spec=ArraySpec( interpolatable=False, dtype=np.uint32, voxel_size=Coordinate((1, 1, 1)), ), )) request = BatchRequest() request.add(img, Coordinate((5, 5, 5))) request.add(swc, Coordinate((5, 5, 5))) with build(pipeline): batch = pipeline.request_batch(request) data = batch[img].data g = batch[swc] assert g.num_vertices() > 0 self.assertNotEqual(data[tuple(np.array(data.shape) // 2)], 0)
def test_voxel_size(self): locations = [[0, 0, 0], [91, 20, 20], [42, 24, 57]] pipeline = ( ExampleSourceSpecifiedLocation(roi=Roi((0, 0, 0), (100, 100, 100)), voxel_size=(5, 2, 2)) + SpecifiedLocation( locations, choose_randomly=False, extra_data=None, jitter=None)) with build(pipeline): batch = pipeline.request_batch( BatchRequest({ ArrayKeys.RAW: ArraySpec(roi=Roi((0, 0, 0), (20, 20, 20))) })) # first locations is skipped # second should start at [80/5, 10/2, 10/2] = [16, 5, 5] self.assertEqual(batch.arrays[ArrayKeys.RAW].data[0, 0, 0], 40255) batch = pipeline.request_batch( BatchRequest({ ArrayKeys.RAW: ArraySpec(roi=Roi((0, 0, 0), (20, 20, 20))) })) # third should start at [30/5, 14/2, 48/2] = [6, 7, 23] self.assertEqual(batch.arrays[ArrayKeys.RAW].data[0, 0, 0], 15374)
def test_precache(self): logging.getLogger("gunpowder.torch.nodes.predict").setLevel( logging.INFO) a = ArrayKey("A") pred = ArrayKey("PRED") model = ExampleModel() reference_request = BatchRequest() reference_request[a] = ArraySpec(roi=Roi((0, 0), (7, 7))) reference_request[pred] = ArraySpec(roi=Roi((1, 1), (5, 5))) source = ExampleTorchTrain2DSource() predict = Predict( model=model, inputs={"a": a}, outputs={0: pred}, array_specs={pred: ArraySpec()}, ) pipeline = source + predict + PreCache(cache_size=3, num_workers=2) request = BatchRequest({ a: ArraySpec(roi=Roi((0, 0), (17, 17))), pred: ArraySpec(roi=Roi((0, 0), (15, 15))), }) # train for a couple of iterations with build(pipeline): batch = pipeline.request_batch(request) assert pred in batch
def test_merge_basics(self): voxel_size = (1, 1, 1) GraphKey("PRESYN") ArrayKey("GT_LABELS") graphsource = GraphTestSource(voxel_size) arraysource = ArrayTestSoure(voxel_size) pipeline = (graphsource, arraysource) + MergeProvider() + RandomLocation() window_request = Coordinate((50, 50, 50)) with build(pipeline): # Check basic merging. request = BatchRequest() request.add((GraphKeys.PRESYN), window_request) request.add((ArrayKeys.GT_LABELS), window_request) batch_res = pipeline.request_batch(request) self.assertTrue(ArrayKeys.GT_LABELS in batch_res.arrays) self.assertTrue(GraphKeys.PRESYN in batch_res.graphs) # Check that request of only one source also works. request = BatchRequest() request.add((GraphKeys.PRESYN), window_request) batch_res = pipeline.request_batch(request) self.assertFalse(ArrayKeys.GT_LABELS in batch_res.arrays) self.assertTrue(GraphKeys.PRESYN in batch_res.graphs) # Check that it fails, when having two sources that provide the same type. arraysource2 = ArrayTestSoure(voxel_size) pipeline_fail = (arraysource, arraysource2) + MergeProvider() + RandomLocation() with self.assertRaises(PipelineSetupError): with build(pipeline_fail): pass
def test_3d(self): test_graph = GraphKey("TEST_GRAPH") graph_spec = GraphSpec(roi=Roi((0, 0, 0), (5, 5, 5))) test_array = ArrayKey("TEST_ARRAY") array_spec = ArraySpec( roi=Roi((0, 0, 0), (5, 5, 5)), voxel_size=Coordinate((1, 1, 1)) ) test_array2 = ArrayKey("TEST_ARRAY2") array2_spec = ArraySpec( roi=Roi((0, 0, 0), (5, 5, 5)), voxel_size=Coordinate((1, 1, 1)) ) snapshot_request = BatchRequest() snapshot_request.add(test_graph, Coordinate((5, 5, 5))) pipeline = ExampleSource( [test_graph, test_array, test_array2], [graph_spec, array_spec, array2_spec] ) + Snapshot( { test_graph: "graphs/graph", test_array: "volumes/array", test_array2: "volumes/array2", }, output_dir=str(self.test_dir), every=2, additional_request=snapshot_request, output_filename="snapshot.hdf", ) snapshot_file_path = Path(self.test_dir, "snapshot.hdf") with build(pipeline): request = BatchRequest() roi = Roi((0, 0, 0), (5, 5, 5)) request[test_array] = ArraySpec(roi=roi) request[test_array2] = ArraySpec(roi=roi) pipeline.request_batch(request) assert snapshot_file_path.exists() f = h5py.File(snapshot_file_path) assert f["volumes/array"] is not None assert f["graphs/graph-ids"] is not None snapshot_file_path.unlink() pipeline.request_batch(request) assert not snapshot_file_path.exists()
def test_without_placeholder(self): test_labels = ArrayKey("TEST_LABELS") test_points = GraphKey("TEST_POINTS") pipeline = ( PointTestSource3D() + RandomLocation(ensure_nonempty=test_points) + ElasticAugment([10, 10, 10], [0.1, 0.1, 0.1], [0, 2.0 * math.pi]) + Snapshot( {test_labels: "volumes/labels"}, output_dir=self.path_to(), output_filename="elastic_augment_test{id}-{iteration}.hdf", )) with build(pipeline): for i in range(2): request_size = Coordinate((40, 40, 40)) request_a = BatchRequest(random_seed=i) request_a.add(test_points, request_size) request_b = BatchRequest(random_seed=i) request_b.add(test_points, request_size) request_b.add(test_labels, request_size) # No array to provide a voxel size to ElasticAugment with pytest.raises(PipelineRequestError): pipeline.request_batch(request_a) batch_b = pipeline.request_batch(request_b) self.assertIn(test_labels, batch_b)
def test_embedding_pipeline( tmpdir, aux_task, blend_mode, fusion_pipeline, train_embedding, snapshot_every ): setup_config = DEFAULT_CONFIG setup_config["FUSION_PIPELINE"] = fusion_pipeline setup_config["TRAIN_EMBEDDING"] = train_embedding setup_config["SNAPSHOT_EVERY"] = snapshot_every setup_config["TENSORBOARD_LOG_DIR"] = tmpdir setup_config["SNAPSHOT_DIR"] = tmpdir setup_config["SNAPSHOT_FILE_NAME"] = "test_snapshot" setup_config["MATCHING_FAILURES_DIR"] = None setup_config["BLEND_MODE"] = blend_mode setup_config["AUX_TASK"] = aux_task voxel_size = Coordinate(setup_config["VOXEL_SIZE"]) output_size = Coordinate(setup_config["OUTPUT_SHAPE"]) * voxel_size input_size = Coordinate(setup_config["INPUT_SHAPE"]) * voxel_size pipeline, raw, output, inputs = embedding_pipeline( setup_config, get_test_data_sources ) request = BatchRequest() request.add(raw, input_size) request.add(output, output_size) for key in inputs: request.add(key, output_size) with build(pipeline): batch = pipeline.request_batch(request) assert output in batch assert raw in batch
def process_function(): scheduler = Client() worker_id = scheduler.context.worker_id num_workers = scheduler.context.num_workers gpu = actor_id_to_gpu_mapping(worker_id) os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu) _logger.info("Worker %d uses gpu %d with %d workers", worker_id, gpu, num_workers) _logger.info("Environment:") for name in os.environ.keys(): _logger.info(' %s:%s', name, os.environ[name]) # _logger.info("GPU is available: %s", tf.test.is_gpu_available()) # from tensorflow.python.client import device_lib # def get_available_gpus(): # local_device_protos = device_lib.list_local_devices() # return [x.name for x in local_device_protos if x.device_type == 'GPU'] # available_gpus = get_available_gpus() # for gpu in available_gpus: # print("Worker %d sees gpus %s" % (actor_id, available_gpus)) import tensorflow as tf with tf.device('/gpu:%d' % 0): from gunpowder import ArrayKey, ArraySpec, build, BatchRequest, DaisyRequestBlocks _RAW = ArrayKey('RAW') roi_map = { ArrayKey('OUTPUT_%d' % i): 'write_roi' for i in range(len(outputs)) } roi_map[_RAW] = 'read_roi' reference = BatchRequest() reference[_RAW] = ArraySpec(roi=None, voxel_size=input_voxel_size) for i in range(len(outputs)): reference[ArrayKey('OUTPUT_%d' % i)] = ArraySpec( roi=None, voxel_size=output_voxel_size) pipeline = pipeline_factory() pipeline += DaisyRequestBlocks(reference=reference, roi_map=roi_map, num_workers=num_cpu_workers) with build(pipeline): pipeline.request_batch(BatchRequest())
def prepare(self, request): deps = BatchRequest() if self.descriptor in request: dims = len(request[self.descriptor].roi.get_shape()) if dims == 2: self.context = self.context[0:2] # increase segmentation ROI to fit Gaussian context_roi = request[self.descriptor].roi.grow( self.context, self.context) # ensure context roi is multiple of voxel size context_roi = context_roi.snap_to_grid(self.voxel_size, mode='shrink') grown_roi = request[self.segmentation].roi.union(context_roi) deps[self.segmentation] = request[self.descriptor].copy() deps[self.segmentation].roi = grown_roi else: self.skip = True if self.unlabelled: deps[self.unlabelled] = deps[self.segmentation].copy() if self.labels_mask: deps[self.labels_mask] = deps[self.segmentation].copy() return deps
def test_relabel_components(self): path = Path(self.path_to("test_swc_source.swc")) # write test swc self._write_swc(path, self._toy_swc_points().to_nx_graph()) # read arrays swc = GraphKey("SWC") source = SwcFileSource(path, [swc]) with build(source): batch = source.request_batch( BatchRequest({swc: GraphSpec(roi=Roi((0, 1, 5), (11, 10, 1)))})) temp_g = batch.points[swc] temp_g.relabel_connected_components() previous_label = None ccs = list(temp_g.connected_components) self.assertEqual(len(ccs), 3) for cc in ccs: self.assertEqual(len(cc), 10) label = None for point_id in cc: if label is None: label = temp_g.node(point_id).attrs["component"] self.assertNotEqual(label, previous_label) self.assertEqual( temp_g.node(point_id).attrs["component"], label) previous_label = label
def prepare(self, request): deps = BatchRequest() deps[self.graph] = request[self.graph].copy() assert (request[self.graph].roi.get_shape() == self.read_size ), f"Got wrong size graph in request" return deps
def test_6_neighborhood(): # array keys graph = GraphKey("GRAPH") neighborhood = ArrayKey("NEIGHBORHOOD") neighborhood_mask = ArrayKey("NEIGHBORHOOD_MASK") distance = 1 pipeline = TestSource(graph) + Neighborhood( graph, neighborhood, neighborhood_mask, distance, array_specs={ neighborhood: ArraySpec(voxel_size=Coordinate((1, 1, 1))), neighborhood_mask: ArraySpec(voxel_size=Coordinate((1, 1, 1))), }, k=6, ) request = BatchRequest() request[neighborhood] = ArraySpec(roi=Roi((0, 0, 0), (10, 10, 10))) request[neighborhood_mask] = ArraySpec(roi=Roi((0, 0, 0), (10, 10, 10))) with build(pipeline): batch = pipeline.request_batch(request) n_data = batch[neighborhood].data n_mask = batch[neighborhood_mask].data masked_ind = list( set([(0, i, 0) for i in range(10) if i not in [0, 4]] + [(i, 5, 0) for i in range(10)] + [(i, 4, 0) for i in range(10) if i not in [0]])) assert all(n_mask[tuple(zip(*masked_ind))] ), f"expected {masked_ind} but saw {np.where(n_mask==1)}"
def run_augmentations( data_providers, roi, keys=(), augmentations=(), voxel_size=lambda key: None): request = BatchRequest() for key in keys: request[key] = ArraySpec(roi(key).snap_to_grid(voxel_size(key)), voxel_size=voxel_size(key)) logger.debug('Requesting batch with request %s', request) data_sources = tuple(provider for provider in data_providers) snapshot = SnapshotAsDict() pipeline = data_sources + RandomProvider() + snapshot for augmentation in augmentations: pipeline += augmentation with build(pipeline) as b: logging.info("submitting request %s", request) batch = b.request_batch(request) logger.debug("Got snapshots from request %s: %s", request, snapshot.snapshots) return batch, snapshot.snapshots[0]
def test_relabel_components(self): path = Path(self.path_to("test_swc_source.swc")) # write test swc self._write_swc(path, self._toy_swc_points()) # read arrays swc = PointsKey("SWC") source = SwcFileSource(path, swc) with build(source): batch = source.request_batch( BatchRequest( {swc: PointsSpec(roi=Roi((0, 1, 0), (11, 10, 1)))})) temp_g = nx.DiGraph() for point_id, point in batch.points[swc].data.items(): temp_g.add_node(point.point_id, label_id=point.label_id) if (point.parent_id != -1 and point.parent_id != point.point_id and point.parent_id in batch.points[swc].data): temp_g.add_edge(point.point_id, point.parent_id) previous_label = None ccs = list(nx.weakly_connected_components(temp_g)) self.assertEqual(len(ccs), 3) for cc in ccs: self.assertEqual(len(cc), 10) label = None for point_id in cc: if label is None: label = temp_g.nodes[point_id]["label_id"] self.assertNotEqual(label, previous_label) self.assertEqual(temp_g.nodes[point_id]["label_id"], label) previous_label = label
def test_with_edge(self): graph_with_edge = GraphKey("TEST_GRAPH_WITH_EDGE") array_with_edge = ArrayKey("RASTERIZED_EDGE") pipeline = GraphTestSourceWithEdge() + RasterizeGraph( GraphKeys.TEST_GRAPH_WITH_EDGE, ArrayKeys.RASTERIZED_EDGE, ArraySpec(voxel_size=(1, 1, 1)), settings=RasterizationSettings(0.5), ) with build(pipeline): request = BatchRequest() roi = Roi((0, 0, 0), (10, 10, 10)) request[GraphKeys.TEST_GRAPH_WITH_EDGE] = GraphSpec(roi=roi) request[ArrayKeys.RASTERIZED_EDGE] = ArraySpec(roi=roi) batch = pipeline.request_batch(request) rasterized = batch.arrays[ArrayKeys.RASTERIZED_EDGE].data assert ( rasterized.sum() == 10 ), f"rasterized has ones at: {np.where(rasterized==1)}"
def test_mirror(self): test_graph = GraphKey("TEST_GRAPH") pipeline = TestSource() + SimpleAugment( mirror_only=[0, 1, 2], transpose_only=[] ) request = BatchRequest() request[GraphKeys.TEST_GRAPH] = GraphSpec(roi=Roi((0, 20, 33), (100, 100, 120))) possible_loc = [[50, 49], [70, 29], [100, 86]] with build(pipeline): seen_mirrored = False for i in range(100): batch = pipeline.request_batch(request) assert len(list(batch[GraphKeys.TEST_GRAPH].nodes)) == 1 node = list(batch[GraphKeys.TEST_GRAPH].nodes)[0] logging.debug(node.location) assert all( [ node.location[dim] in possible_loc[dim] for dim in range(3) ] ) seen_mirrored = seen_mirrored or any( [node.location[dim] == possible_loc[dim][1] for dim in range(3)] ) assert Roi((0, 20, 33), (100, 100, 120)).contains(batch[GraphKeys.TEST_GRAPH].spec.roi) assert batch[GraphKeys.TEST_GRAPH].spec.roi.contains(node.location) assert seen_mirrored
def test_square(self): test_graph = GraphKey("TEST_GRAPH") test_array1 = ArrayKey("TEST_ARRAY1") test_array2 = ArrayKey("TEST_ARRAY2") pipeline = ((ArrayTestSource(), TestSource()) + MergeProvider() + Pad(test_array1, None) + Pad(test_array2, None) + Pad(test_graph, None) + SimpleAugment( mirror_only=[1,2], transpose_only=[1,2] )) request = BatchRequest() request[GraphKeys.TEST_GRAPH] = GraphSpec(roi=Roi((0, 50, 65), (100, 100, 100))) request[ArrayKeys.TEST_ARRAY1] = ArraySpec(roi=Roi((0, 0, 15), (100, 200, 200))) request[ArrayKeys.TEST_ARRAY2] = ArraySpec(roi=Roi((0, 50, 65), (100, 100, 100))) with build(pipeline): for i in range(100): batch = pipeline.request_batch(request) assert len(list(batch[GraphKeys.TEST_GRAPH].nodes)) == 1 for (array_key, array) in batch.arrays.items(): assert batch.arrays[array_key].data.shape == batch.arrays[array_key].spec.roi.get_shape()
def prepare(self, request): deps = BatchRequest() deps[self.dense_mst] = request[self.mst].copy() deps[self.mst] = request[self.mst].copy() deps[self.embeddings] = ArraySpec(roi=request[self.mst].roi) return deps
def prepare(self, request): deps = BatchRequest() deps[self.msts[0]] = request[self.msts[1]].copy() if (self.msts_dense is not None) and (self.msts_dense[1] in request): deps[self.msts_dense[0]] = request[self.msts_dense[1]].copy() return deps
def test_prepare1(self): key = ArrayKey("TEST_ARRAY") spec = ArraySpec(voxel_size=Coordinate((1, 1)), interpolatable=True) hdf5_source = Hdf5Source(self.fake_data_file, {key: "testdata"}, array_specs={key: spec}) request = BatchRequest() shape = Coordinate((3, 3)) request.add(key, shape, voxel_size=Coordinate((1, 1))) shift_node = ShiftAugment(sigma=1, shift_axis=0) with build((hdf5_source + shift_node)): shift_node.prepare(request) self.assertTrue(shift_node.ndim == 2) self.assertTrue(shift_node.shift_sigmas == tuple([0.0, 1.0]))
def prepare(self, request): assert ArrayKeys.C in request dependencies = BatchRequest() dependencies[ArrayKeys.B] = ArraySpec(request[ArrayKeys.C].roi.grow( self.context, self.context)) return dependencies
def prepare(self, request: BatchRequest): deps = BatchRequest() if self.output is not None: deps[self.array] = request[self.output].copy() else: deps[self.array] = request[self.array].copy() return deps
def test_pipeline2(self): key = ArrayKey("TEST_ARRAY") spec = ArraySpec(voxel_size=Coordinate((3, 1)), interpolatable=True) hdf5_source = Hdf5Source(self.fake_data_file, {key: "testdata"}, array_specs={key: spec}) request = BatchRequest() shape = Coordinate((3, 3)) request.add(key, shape, voxel_size=Coordinate((3, 1))) shift_node = ShiftAugment(prob_slip=0.2, prob_shift=0.2, sigma=1, shift_axis=0) with build((hdf5_source + shift_node)) as b: b.request_batch(request)
def prepare(self, request: BatchRequest, seed: int, direction: Coordinate) -> Tuple[BatchRequest, int]: """ Only request everything with the given seed """ dps = BatchRequest(random_seed=seed) if self.nonempty_placeholder is not None: # request nonempty placeholder of size request total roi # grow such that it can be cropped down to two different locations growth = self._get_growth() total_roi = request.get_total_roi() grown_roi = total_roi.grow(growth, growth) dps[self.nonempty_placeholder] = GraphSpec(roi=grown_roi, placeholder=True) # handle smaller requests array_keys = list(request.array_specs.keys()) voxel_size = self.spec.get_lcm_voxel_size(array_keys) direction = Coordinate(direction) direction -= Coordinate( tuple(np.array(direction) % np.array(voxel_size))) if any([points in request for points in self.points]): dps[self.point_source] = copy.deepcopy(request[self.points[0]]) dps[self.point_source].roi = dps[self.point_source].roi.shift( direction) if any([array in request for array in self.arrays]): dps[self.array_source] = copy.deepcopy(request[self.arrays[0]]) dps[self.array_source].roi = dps[self.array_source].roi.shift( direction) if any([labels in request for labels in self.labels]): dps[self.label_source] = copy.deepcopy(request[self.labels[0]]) dps[self.label_source].roi = dps[self.label_source].roi.shift( direction) for source, targets in self.extra_keys.items(): if targets[0] in request: dps[source] = copy.deepcopy(request[targets[0]]) dps[source].roi = dps[source].roi.shift(direction) return dps
def test_transpose(): voxel_size = Coordinate((20, 20)) graph_key = GraphKey("GRAPH") array_key = ArrayKey("ARRAY") graph = Graph( [Node(id=1, location=np.array([450, 550]))], [], GraphSpec(roi=Roi((100, 200), (800, 600))), ) data = np.zeros([40, 30]) data[17, 17] = 1 array = Array( data, ArraySpec(roi=Roi((100, 200), (800, 600)), voxel_size=voxel_size)) default_pipeline = ( (GraphSource(graph_key, graph), ArraySource(array_key, array)) + MergeProvider() + SimpleAugment( mirror_only=[], transpose_only=[0, 1], transpose_probs=[0, 0])) transpose_pipeline = ( (GraphSource(graph_key, graph), ArraySource(array_key, array)) + MergeProvider() + SimpleAugment( mirror_only=[], transpose_only=[0, 1], transpose_probs=[1, 1])) request = BatchRequest() request[graph_key] = GraphSpec(roi=Roi((400, 500), (200, 300))) request[array_key] = ArraySpec(roi=Roi((400, 500), (200, 300))) with build(default_pipeline): expected_location = [450, 550] batch = default_pipeline.request_batch(request) assert len(list(batch[graph_key].nodes)) == 1 node = list(batch[graph_key].nodes)[0] assert all(np.isclose(node.location, expected_location)) node_voxel_index = Coordinate( (node.location - batch[array_key].spec.roi.get_offset()) / voxel_size) assert ( batch[array_key].data[node_voxel_index] == 1 ), f"Node at {np.where(batch[array_key].data == 1)} not {node_voxel_index}" with build(transpose_pipeline): expected_location = [410, 590] batch = transpose_pipeline.request_batch(request) assert len(list(batch[graph_key].nodes)) == 1 node = list(batch[graph_key].nodes)[0] assert all(np.isclose(node.location, expected_location)) node_voxel_index = Coordinate( (node.location - batch[array_key].spec.roi.get_offset()) / voxel_size) assert ( batch[array_key].data[node_voxel_index] == 1 ), f"Node at {np.where(batch[array_key].data == 1)} not {node_voxel_index}"
def test_3d_basics(self): test_labels = ArrayKey("TEST_LABELS") test_points = GraphKey("TEST_POINTS") test_raster = ArrayKey("TEST_RASTER") pipeline = ( PointTestSource3D() + ElasticAugment( [10, 10, 10], [0.1, 0.1, 0.1], # [0, 0, 0], # no jitter [0, 2.0 * math.pi], ) + RasterizeGraph( test_points, test_raster, settings=RasterizationSettings(radius=2, mode="peak"), ) + Snapshot( { test_labels: "volumes/labels", test_raster: "volumes/raster" }, dataset_dtypes={test_raster: np.float32}, output_dir=self.path_to(), output_filename="elastic_augment_test{id}-{iteration}.hdf", )) for _ in range(5): with build(pipeline): request_roi = Roi((-20, -20, -20), (40, 40, 40)) request = BatchRequest() request[test_labels] = ArraySpec(roi=request_roi) request[test_points] = GraphSpec(roi=request_roi) request[test_raster] = ArraySpec(roi=request_roi) batch = pipeline.request_batch(request) labels = batch[test_labels] points = batch[test_points] # the point at (0, 0, 0) should not have moved self.assertTrue(points.contains(0)) labels_data_roi = ( labels.spec.roi - labels.spec.roi.get_begin()) / labels.spec.voxel_size # points should have moved together with the voxels for point in points.nodes: loc = point.location - labels.spec.roi.get_begin() loc = loc / labels.spec.voxel_size loc = Coordinate(int(round(x)) for x in loc) if labels_data_roi.contains(loc): self.assertEqual(labels.data[loc], point.id)
def test_realistic_valid_examples(example, use_gurobi): penalty_attr = "penalty" location_attr = "location" example_dir = Path(__file__).parent / "mouselight_examples" / "valid" / example consensus = PointsKey("CONSENSUS") skeletonization = PointsKey("SKELETONIZATION") matched = PointsKey("MATCHED") matched_with_fallback = PointsKey("MATCHED_WITH_FALLBACK") inf_roi = Roi(Coordinate((None,) * 3), Coordinate((None,) * 3)) request = BatchRequest() request[matched] = PointsSpec(roi=inf_roi) request[matched_with_fallback] = PointsSpec(roi=inf_roi) request[consensus] = PointsSpec(roi=inf_roi) pipeline = ( ( GraphSource(example_dir / "graph.obj", [skeletonization]), GraphSource(example_dir / "tree.obj", [consensus]), ) + MergeProvider() + TopologicalMatcher( skeletonization, consensus, matched, expected_edge_len=10, match_distance_threshold=76, max_gap_crossing=48, use_gurobi=use_gurobi, location_attr=location_attr, penalty_attr=penalty_attr, ) + TopologicalMatcher( skeletonization, consensus, matched_with_fallback, expected_edge_len=10, match_distance_threshold=76, max_gap_crossing=48, use_gurobi=use_gurobi, location_attr=location_attr, penalty_attr=penalty_attr, with_fallback=True, ) ) with build(pipeline): batch = pipeline.request_batch(request) consensus_ccs = list(batch[consensus].connected_components) matched_with_fallback_ccs = list(batch[matched_with_fallback].connected_components) matched_ccs = list(batch[matched].connected_components) assert len(matched_ccs) == len(consensus_ccs)
def test_output(self): meta_base = self.path_to('tf_graph') ArrayKey('A') ArrayKey('B') ArrayKey('C') ArrayKey('GRADIENT_A') # create model meta graph file and get input/output names (a, b, c, optimizer, loss) = self.create_meta_graph(meta_base) source = ExampleTensorflowTrainSource() train = Train( meta_base, optimizer=optimizer, loss=loss, inputs={a: ArrayKeys.A, b: ArrayKeys.B}, outputs={c: ArrayKeys.C}, gradients={a: ArrayKeys.GRADIENT_A}, save_every=100) pipeline = source + train request = BatchRequest({ ArrayKeys.A: ArraySpec(roi=Roi((0, 0), (2, 2))), ArrayKeys.B: ArraySpec(roi=Roi((0, 0), (2, 2))), ArrayKeys.C: ArraySpec(roi=Roi((0, 0), (2, 2))), ArrayKeys.GRADIENT_A: ArraySpec(roi=Roi((0, 0), (2, 2))), }) # train for a couple of iterations with build(pipeline): batch = pipeline.request_batch(request) self.assertAlmostEqual(batch.loss, 9.8994951) gradient_a = batch.arrays[ArrayKeys.GRADIENT_A].data self.assertTrue(gradient_a[0, 0] < gradient_a[0, 1]) self.assertTrue(gradient_a[0, 1] < gradient_a[1, 0]) self.assertTrue(gradient_a[1, 0] < gradient_a[1, 1]) for i in range(200-1): loss1 = batch.loss batch = pipeline.request_batch(request) loss2 = batch.loss self.assertLess(loss2, loss1) # resume training with build(pipeline): for i in range(100): loss1 = batch.loss batch = pipeline.request_batch(request) loss2 = batch.loss self.assertLess(loss2, loss1)