soorena62
58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
from radiuma_api.workflow import Workflow
|
|
from radiuma_api.execution_context import ExecutionContext
|
|
from radiuma_api.io_port import InPort, OutPort, ImageType, TextType
|
|
from radiuma_api.module import Module
|
|
from radiuma_api.task import Status
|
|
from radiuma_api.asset import Asset
|
|
# Codes Go below:
|
|
|
|
|
|
class Reader(Module):
|
|
def __init__(self):
|
|
super().__init__("Reader")
|
|
self._add_output_port(OutPort("images", ImageType()))
|
|
|
|
def run(self, context):
|
|
context.put(self.outputs[0], Asset(["img1"], "images"))
|
|
self._set_status(Status.COMPLETED)
|
|
|
|
class Extractor(Module):
|
|
def __init__(self):
|
|
super().__init__("Extractor")
|
|
self._add_input_port(InPort("images", ImageType()))
|
|
self._add_output_port(OutPort("features", TextType()))
|
|
|
|
def run(self, context):
|
|
imgs = context.get(self.inputs[0]).data
|
|
context.put(self.outputs[0], Asset([{"img": imgs[0]}], "features"))
|
|
self._set_status(Status.COMPLETED)
|
|
|
|
class Writer(Module):
|
|
def __init__(self):
|
|
super().__init__("Writer")
|
|
self._add_input_port(InPort("features", TextType()))
|
|
|
|
def run(self, context):
|
|
feats = context.get(self.inputs[0]).data
|
|
assert feats[0]["img"] == "img1"
|
|
self._set_status(Status.COMPLETED)
|
|
|
|
def test_end_to_end():
|
|
ctx = ExecutionContext()
|
|
wf = Workflow("WF")
|
|
|
|
r = Reader()
|
|
e = Extractor()
|
|
w = Writer()
|
|
|
|
r.outputs[0].connect(e.inputs[0])
|
|
e.outputs[0].connect(w.inputs[0])
|
|
|
|
wf.add_child(r)
|
|
wf.add_child(e)
|
|
wf.add_child(w)
|
|
|
|
wf.run(ctx)
|
|
|
|
assert w.status == Status.COMPLETED
|