class ExampleGraphSource(BatchProvider):
def __init__(self):
self.dtype = float
self.__vertices = [
Node(id=1, location=np.array([1, 1, 1], dtype=self.dtype)),
Node(id=2, location=np.array([500, 500, 500], dtype=self.dtype)),
Node(id=3, location=np.array([550, 550, 550], dtype=self.dtype)),
]
self.__edges = [Edge(1, 2), Edge(2, 3)]
self.__spec = GraphSpec(roi=Roi(Coordinate([-500, -500, -500]),
Coordinate([1500, 1500, 1500])))
self.graph = Graph(self.__vertices, self.__edges, self.__spec)
def setup(self):
self.provides(GraphKeys.TEST_GRAPH, self.__spec)
def provide(self, request):
batch = Batch()
roi = request[GraphKeys.TEST_GRAPH].roi
sub_graph = self.graph.crop(roi)
batch[GraphKeys.TEST_GRAPH] = sub_graph
return batch
class TestPointSource(BatchProvider):
def __init__(self, points: List[GraphKey], directed: bool,
size: Coordinate, num_points: int):
self.points = points
self.directed = directed
self.size = size
self.num_points = num_points
def setup(self):
roi = Roi(Coordinate([0] * len(self.size)), self.size)
for points_key in self.points:
self.provides(points_key, GraphSpec(roi=roi,
directed=self.directed))
k = min(self.size)
nodes = [
Node(id=i, location=np.array([i * k / self.num_points] * 3))
for i in range(self.num_points)
]
edges = [Edge(i, i + 1) for i in range(self.num_points - 1)]
self.graph = Graph(nodes, edges,
GraphSpec(roi=roi, directed=self.directed))
def provide(self, request: BatchRequest) -> Batch:
batch = Batch()
for points_key in self.points:
if points_key in request:
spec = request[points_key].copy()
subgraph = self.graph.crop(roi=spec.roi)
subgraph.relabel_connected_components()
batch[points_key] = subgraph
return batch
class TestSourceRandomLocation(BatchProvider):
def __init__(self):
self.graph = Graph(
[
Node(id=1, location=np.array([1, 1, 1])),
Node(id=2, location=np.array([500, 500, 500])),
Node(id=3, location=np.array([550, 550, 550])),
],
[],
GraphSpec(roi=Roi((0, 0, 0), (1000, 1000, 1000))),
)
def setup(self):
self.provides(GraphKeys.TEST_GRAPH, self.graph.spec)
def provide(self, request):
batch = Batch()
roi = request[GraphKeys.TEST_GRAPH].roi
batch[GraphKeys.TEST_GRAPH] = self.graph.crop(roi).trim(roi)
return batch
def test_crop(self):
g = Graph(self.nodes, self.edges, self.spec)
sub_g = g.crop(Roi(Coordinate([1, 1, 1]), Coordinate([3, 3, 3])))
self.assertEqual(g.spec.roi, self.spec.roi)
self.assertEqual(sub_g.spec.roi,
Roi(Coordinate([1, 1, 1]), Coordinate([3, 3, 3])))
sub_g.spec.directed = False
self.assertTrue(g.spec.directed)
self.assertFalse(sub_g.spec.directed)
def provide(self, request):
outputs = Batch()
nodes = [
Node(id=0, location=np.array((1, 1, 1))),
Node(id=1, location=np.array((10, 10, 10))),
Node(id=2, location=np.array((19, 19, 19))),
Node(id=3, location=np.array((21, 21, 21))),
Node(id=104, location=np.array((30, 30, 30))),
Node(id=5, location=np.array((39, 39, 39))),
]
edges = [Edge(0, 1), Edge(1, 2), Edge(3, 104), Edge(104, 5)]
spec = self.spec[GraphKeys.RAW].copy()
spec.roi = request[GraphKeys.RAW].roi
graph = Graph(nodes, edges, spec)
outputs[GraphKeys.RAW] = graph.crop(spec.roi)
return outputs
class ExampleSourceRandomLocation(BatchProvider):
def __init__(self):
self.graph = Graph([
Node(1, np.array([1, 1, 1])),
Node(2, np.array([500, 500, 500])),
Node(3, np.array([550, 550, 550])),
], [], GraphSpec(roi=Roi((-500, -500, -500), (1500, 1500, 1500))))
def setup(self):
self.provides(GraphKeys.TEST_POINTS, self.graph.spec)
def provide(self, request):
batch = Batch()
roi = request[GraphKeys.TEST_POINTS].roi
batch[GraphKeys.TEST_POINTS] = self.graph.crop(roi).trim(roi)
return batch
class TestSource(BatchProvider):
def __init__(self):
self.graph = Graph(
[Node(id=1, location=np.array([50, 70, 100]))],
[],
GraphSpec(roi=Roi((-200, -200, -200), (400, 400, 478))),
)
def setup(self):
self.provides(GraphKeys.TEST_GRAPH, self.graph.spec)
def prepare(self, request):
return request
def provide(self, request):
batch = Batch()
roi = request[GraphKeys.TEST_GRAPH].roi
batch[GraphKeys.TEST_GRAPH] = self.graph.crop(roi).trim(roi)
return batch