class GraphTestSourceWithEdge(BatchProvider):
def __init__(self):
self.voxel_size = Coordinate((1, 1, 1))
self.nodes = [
# corners
Node(id=1, location=np.array((0, 4, 4))),
Node(id=2, location=np.array((9, 4, 4)))
]
self.edges = [
Edge(1, 2)
]
self.graph_spec = GraphSpec(roi=Roi((0, 0, 0), (10, 10, 10)))
self.graph = Graph(self.nodes, self.edges, self.graph_spec)
def setup(self):
self.provides(
GraphKeys.TEST_GRAPH_WITH_EDGE,
self.graph_spec,
)
def provide(self, request):
batch = Batch()
graph_roi = request[GraphKeys.TEST_GRAPH_WITH_EDGE].roi
batch.graphs[GraphKeys.TEST_GRAPH_WITH_EDGE] = self.graph.crop(graph_roi).trim(
graph_roi
)
return batch
class GraphTestSource3D(BatchProvider):
def __init__(self):
self.voxel_size = Coordinate((40, 4, 4))
self.nodes = [
# corners
Node(id=1, location=np.array((-200, -200, -200))),
Node(id=2, location=np.array((-200, -200, 199))),
Node(id=3, location=np.array((-200, 199, -200))),
Node(id=4, location=np.array((-200, 199, 199))),
Node(id=5, location=np.array((199, -200, -200))),
Node(id=6, location=np.array((199, -200, 199))),
Node(id=7, location=np.array((199, 199, -200))),
Node(id=8, location=np.array((199, 199, 199))),
# center
Node(id=9, location=np.array((0, 0, 0))),
Node(id=10, location=np.array((-1, -1, -1))),
]
self.graph_spec = GraphSpec(roi=Roi((-100, -100, -100), (300, 300, 300)))
self.array_spec = ArraySpec(
roi=Roi((-200, -200, -200), (400, 400, 400)), voxel_size=self.voxel_size
)
self.graph = Graph(self.nodes, [], self.graph_spec)
def setup(self):
self.provides(
GraphKeys.TEST_GRAPH,
self.graph_spec,
)
self.provides(
ArrayKeys.GT_LABELS,
self.array_spec,
)
def provide(self, request):
batch = Batch()
graph_roi = request[GraphKeys.TEST_GRAPH].roi
batch.graphs[GraphKeys.TEST_GRAPH] = self.graph.crop(graph_roi).trim(graph_roi)
roi_array = request[ArrayKeys.GT_LABELS].roi
image = np.ones(roi_array.get_shape() / self.voxel_size, dtype=np.uint64)
# label half of GT_LABELS differently
depth = image.shape[0]
image[0 : depth // 2] = 2
spec = self.spec[ArrayKeys.GT_LABELS].copy()
spec.roi = roi_array
batch.arrays[ArrayKeys.GT_LABELS] = Array(image, spec=spec)
return batch