269 lines
10 KiB
Python
269 lines
10 KiB
Python
"""Utility classes to help with the simpler Python interface
|
|
for connecting and executing pipelines."""
|
|
|
|
__all__ = ['select_ports', 'Pipeline', 'Output']
|
|
|
|
def _call(first, last, inp=None, port=0):
|
|
"""Set the input of the first filter, update the pipeline
|
|
and return the output."""
|
|
if inp and not first.GetNumberOfInputPorts():
|
|
raise ValueError(f"{first.GetClassName()} does not have input ports yet an input was passed to the pipeline.")
|
|
in_cons = []
|
|
if first.GetNumberOfInputPorts():
|
|
n_cons = first.GetNumberOfInputConnections(port)
|
|
for i in range(n_cons):
|
|
op = first.GetInputConnection(port, i)
|
|
if op and op.GetProducer():
|
|
op.GetProducer().Register(None)
|
|
in_cons.append(op)
|
|
from vtkmodules.vtkCommonExecutionModel import vtkAlgorithm
|
|
from vtkmodules.vtkCommonExecutionModel import vtkTrivialProducer
|
|
from collections.abc import Sequence
|
|
if isinstance(inp, Sequence):
|
|
if first.GetInputPortInformation(port).Has(
|
|
vtkAlgorithm.INPUT_IS_REPEATABLE()):
|
|
first.RemoveAllInputConnections(port)
|
|
for aInp in inp:
|
|
tp = vtkTrivialProducer()
|
|
tp.SetOutput(aInp)
|
|
first.AddInputConnection(port, tp.GetOutputPort());
|
|
else:
|
|
tp = vtkTrivialProducer()
|
|
tp.SetOutput(inp)
|
|
first.SetInputConnection(port, tp.GetOutputPort());
|
|
|
|
output = last.update().output
|
|
|
|
if first.GetNumberOfInputPorts():
|
|
first.RemoveAllInputConnections(port)
|
|
for op in in_cons:
|
|
first.AddInputConnection(port, op)
|
|
if op and op.GetProducer():
|
|
op.GetProducer().UnRegister(None)
|
|
|
|
output_copy = []
|
|
if type(output) is not tuple:
|
|
output = (output,)
|
|
for opt in output:
|
|
copy = opt.NewInstance()
|
|
copy.ShallowCopy(opt)
|
|
output_copy.append(copy)
|
|
if len(output_copy) == 1:
|
|
return output_copy[0]
|
|
else:
|
|
return tuple(output_copy)
|
|
|
|
|
|
class select_ports(object):
|
|
"""Helper class for selecting input and output ports when
|
|
connecting pipeline objects with the >> operator.
|
|
Example uses:
|
|
# Connect a source to the second input of a filter.
|
|
source >> select_ports(1, filter)
|
|
# Connect the second output of a source to a filter.
|
|
select_ports(source, 1) >> filter
|
|
# Combination of both: Connect source to second
|
|
# input of the filter, then connect the second
|
|
# output of that filter to another one.
|
|
source >>> select_ports(1, filter, 1) >> filter2
|
|
"""
|
|
def __init__(self, *args):
|
|
"""This constructor takes 2 or 3 arguments.
|
|
The possibilities are:
|
|
select_ports(input_port, algorithm)
|
|
select_ports(algorithm, output_port)
|
|
select_ports(input_port, algorithm, output_port)
|
|
"""
|
|
nargs = len(args)
|
|
if nargs < 2 or nargs > 3:
|
|
raise ValueError("Expecting 2 or 3 arguments")
|
|
|
|
self.input_port = None
|
|
self.output_port = None
|
|
before_alg = True
|
|
for arg in args:
|
|
if hasattr(arg, "IsA") and arg.IsA("vtkAlgorithm"):
|
|
self.algorithm = arg
|
|
before_alg = False
|
|
else:
|
|
if before_alg:
|
|
self.input_port = arg
|
|
else:
|
|
self.output_port = arg
|
|
if not self.input_port:
|
|
self.input_port = 0
|
|
if not self.output_port:
|
|
self.output_port = 0
|
|
|
|
def SetInputConnection(self, inp):
|
|
"Forwards to underlying algorithm and port."
|
|
self.algorithm.SetInputConnection(self.input_port, inp)
|
|
|
|
def AddInputConnection(self, inp):
|
|
"Forwards to underlying algorithm and port."
|
|
self.algorithm.AddInputConnection(self.input_port, inp)
|
|
|
|
def GetOutputPort(self):
|
|
"Returns the output port of the underlying algorithm."
|
|
return self.algorithm.GetOutputPort(self.output_port)
|
|
|
|
def GetInputPortInformation(self, port):
|
|
return self.algorithm.GetInputPortInformation(self.input_port)
|
|
|
|
def update(self):
|
|
"""Execute the algorithm and return the output from the selected
|
|
output port."""
|
|
return self.algorithm.update()
|
|
|
|
def __rshift__(self, rhs):
|
|
"Creates a pipeline between the underlying port and an algorithm."
|
|
return Pipeline(self, rhs)
|
|
|
|
def __rrshift__(self, lhs):
|
|
"""Creates a pipeline between the underlying port and an algorithm.
|
|
This is to handle sequence >> select_ports where the port can
|
|
accept multiple connections."""
|
|
from collections.abc import Sequence
|
|
if lhs is None or (isinstance(lhs, Sequence) and len(lhs == 0)):
|
|
self.algorithm.RemoveAllInputConnections(self.input_port)
|
|
return self
|
|
return Pipeline(lhs, self)
|
|
|
|
def __call__(self, inp=None):
|
|
"""Executes the underlying algorithm by passing input data to
|
|
the selected input port. Returns a single output or a tuple
|
|
if there are multiple outputs."""
|
|
return _call(self.algorithm, self.algorithm, inp, self.input_port)
|
|
|
|
class Pipeline(object):
|
|
"""Pipeline objects are created when 2 or more algorithms are
|
|
connected with the >> operator. They store the first and last
|
|
algorithms in the pipeline and enable connecting more algorithms
|
|
and executing the pipeline. One should not have to create Pipeline
|
|
objects directly. They are created by the use of the >> operator."""
|
|
|
|
PIPELINE = 0
|
|
ALGORITHM = 1
|
|
DATA = 2
|
|
UNKNOWN = 3
|
|
|
|
def __init__(self, lhs, rhs):
|
|
"""Create a pipeline object that connects two objects of the
|
|
following type: data object, pipeline object, algorithm object."""
|
|
left_type = self._determine_type(lhs)
|
|
right_type = self._determine_type(rhs)
|
|
if right_type == Pipeline.ALGORITHM:
|
|
rhs_alg = rhs
|
|
elif right_type == Pipeline.PIPELINE:
|
|
rhs_alg = rhs.first
|
|
else:
|
|
raise TypeError(
|
|
f"unsupported operand type(s) for >>: {type(lhs).__name__} and {type(rhs).__name__}")
|
|
|
|
from collections.abc import Sequence
|
|
if isinstance(lhs, Sequence):
|
|
for inp in lhs:
|
|
self._connect(inp, rhs, rhs_alg, "AddInputConnection")
|
|
else:
|
|
self._connect(lhs, rhs, rhs_alg, "SetInputConnection")
|
|
|
|
def _connect(self, lhs, rhs, rhs_alg, connect_method):
|
|
from vtkmodules.vtkCommonExecutionModel import vtkAlgorithm
|
|
inInfo = rhs_alg.GetInputPortInformation(0)
|
|
if inInfo.Has(vtkAlgorithm.INPUT_IS_REPEATABLE()):
|
|
connect_method = 'AddInputConnection'
|
|
|
|
left_type = self._determine_type(lhs)
|
|
right_type = self._determine_type(rhs)
|
|
if left_type == Pipeline.UNKNOWN:
|
|
raise TypeError(
|
|
f"unsupported operand type(s) for >>: {type(lhs).__name__} and {type(rhs).__name__}")
|
|
if right_type == Pipeline.ALGORITHM:
|
|
if left_type == Pipeline.ALGORITHM:
|
|
getattr(rhs_alg, connect_method)(lhs.GetOutputPort())
|
|
self.first = lhs
|
|
self.last = rhs
|
|
elif left_type == Pipeline.PIPELINE:
|
|
getattr(rhs_alg, connect_method)(lhs.last.GetOutputPort())
|
|
self.first = lhs.first
|
|
self.last = rhs
|
|
elif left_type == Pipeline.DATA:
|
|
from vtkmodules.vtkCommonExecutionModel import vtkTrivialProducer
|
|
source = vtkTrivialProducer()
|
|
source.SetOutput(lhs)
|
|
getattr(rhs_alg, connect_method)(source.GetOutputPort())
|
|
self.first = source
|
|
self.last = rhs
|
|
elif right_type == Pipeline.PIPELINE:
|
|
if left_type == Pipeline.ALGORITHM:
|
|
self.first = lhs
|
|
self.last = rhs.last
|
|
getattr(rhs_alg, connect_method)(lhs.GetOutputPort())
|
|
elif left_type == Pipeline.PIPELINE:
|
|
getattr(rhs_alg, connect_method)(lhs.last.GetOutputPort())
|
|
self.first = lhs.first
|
|
self.last = rhs.last
|
|
elif left_type == Pipeline.DATA:
|
|
from vtkmodules.vtkCommonExecutionModel import vtkTrivialProducer
|
|
source = vtkTrivialProducer()
|
|
source.SetOutput(lhs)
|
|
getattr(rhs_alg, connect_method)(source.GetOutputPort())
|
|
self.first = source
|
|
self.last = rhs.last
|
|
|
|
def _determine_type(self, arg):
|
|
if type(arg) is Pipeline:
|
|
return Pipeline.PIPELINE
|
|
if hasattr(arg, "SetInputConnection"):
|
|
return Pipeline.ALGORITHM
|
|
if hasattr(arg, "IsA") and arg.IsA("vtkDataObject"):
|
|
return Pipeline.DATA
|
|
return Pipeline.UNKNOWN
|
|
|
|
def update(self, **kwargs):
|
|
"""Update the pipeline and return the last algorithm's
|
|
output."""
|
|
return self.last.update()
|
|
|
|
def __call__(self, inp=None):
|
|
"""Sets the input of the first filter, update the pipeline
|
|
and returns the output. A single data object or a tuple
|
|
of data objects (when there are multiple outputs) are
|
|
returned."""
|
|
return _call(self.first, self.last, inp)
|
|
|
|
def __rshift__(self, rhs):
|
|
"""Used to connect two pipeline items. The left side can
|
|
be a data object, an algorithm or a pipeline. The right
|
|
side can be an algorithm or a pipeline."""
|
|
return Pipeline(self, rhs)
|
|
|
|
def __rrshift__(self, lhs):
|
|
"""Creates a pipeline between a sequence input and a pipeline."""
|
|
from collections.abc import Sequence
|
|
if lhs is None or (isinstance(lhs, Sequence) and len(lhs) == 0):
|
|
self.first.RemoveAllInputConnections(0)
|
|
return self
|
|
return Pipeline(lhs, self)
|
|
|
|
class Output(object):
|
|
"""Helper object to represent the output of an algorithms as
|
|
returned by the update() method. Implements the output property
|
|
enabling calling update().output."""
|
|
def __init__(self, algorithm, **kwargs):
|
|
self.algorithm = algorithm
|
|
self.algorithm.Update()
|
|
|
|
@property
|
|
def output(self):
|
|
"""Returns a single data object or a tuple of data objects
|
|
if there are multiple outputs."""
|
|
if self.algorithm.GetNumberOfOutputPorts() == 1:
|
|
return self.algorithm.GetOutputDataObject(0)
|
|
else:
|
|
outputs = []
|
|
nOutputs = self.algorithm.GetNumberOfOutputPorts()
|
|
for i in range(nOutputs):
|
|
outputs.append(self.algorithm.GetOutputDataObject(i))
|
|
return tuple(outputs)
|