40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
import numpy as np
|
|
|
|
from .event_bus import EventBus
|
|
|
|
|
|
class ImagePipeline:
|
|
def __init__(self, bus: EventBus):
|
|
self.bus = bus
|
|
self.stages = []
|
|
self.stagedata = {}
|
|
|
|
def register_stage(self, name: str):
|
|
self.stages.append(name)
|
|
self.stagedata[len(self.stages)-1] = None
|
|
self.bus.publish_deferred("pipeline_stages", self.stages)
|
|
return len(self.stages) - 1
|
|
|
|
def rename_stage(self, id: int, name: str):
|
|
if id >= 0 and id < len(self.stages):
|
|
self.stages[id] = name
|
|
self.bus.publish_deferred("pipeline_stages", self.stages)
|
|
|
|
def publish(self, id: int, img: np.ndarray):
|
|
self.stagedata[id] = img.astype(np.float32)
|
|
self.bus.publish_deferred("pipeline_stage", (id, self.stagedata[id]))
|
|
|
|
def get_stage_data(self, id: int):
|
|
if id >= 0 and id < len(self.stages):
|
|
return self.stagedata[id]
|
|
else:
|
|
return None
|
|
|
|
def get_stage_name(self, id:int):
|
|
if id >= 0 and id < len(self.stages):
|
|
return self.stages[id]
|
|
else:
|
|
return None
|
|
|
|
def republish_stages(self):
|
|
self.bus.publish_deferred("pipeline_stages", self.stages) |