74 lines
2.2 KiB
Python
74 lines
2.2 KiB
Python
import numpy as np
|
|
|
|
from .event_bus import EventBus
|
|
|
|
|
|
class ImagePipeline:
|
|
def __init__(self, bus: EventBus):
|
|
self.bus = bus
|
|
self.id_counter = 0
|
|
self.stages = {}
|
|
self.stagedata = {}
|
|
self.stagedata_full = {}
|
|
|
|
def load_stages(self, stages:dict):
|
|
self.stages = stages
|
|
self.stagedata.clear()
|
|
self.stagedata_full.clear()
|
|
self.id_counter = len(stages)
|
|
for id, stage in self.stages.items():
|
|
print(id, stage)
|
|
self.stagedata[id] = None
|
|
self.stagedata_full[id] = None
|
|
|
|
def register_stage(self, name: str):
|
|
self.stages[self.id_counter] = name
|
|
self.stagedata[self.id_counter] = None
|
|
self.stagedata_full[self.id_counter] = None
|
|
self.bus.publish_deferred("pipeline_stages", self.stages)
|
|
self.id_counter += 1
|
|
return self.id_counter-1
|
|
|
|
def rename_stage(self, id: int, name: str):
|
|
if id in self.stages:
|
|
self.stages[id] = name
|
|
self.bus.publish_deferred("pipeline_stages", self.stages)
|
|
|
|
def publish(self, id: int, img: np.ndarray, full_res=False):
|
|
if img is None:
|
|
return
|
|
if full_res:
|
|
self.stagedata_full[id] = img.astype(np.float32)
|
|
self.bus.publish_deferred(
|
|
"pipeline_stage_full", (id, self.stagedata_full[id]))
|
|
else:
|
|
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 in self.stagedata:
|
|
return self.stagedata[id]
|
|
else:
|
|
return None
|
|
|
|
def get_stage_data_full(self, id: int):
|
|
if id in self.stagedata_full:
|
|
return self.stagedata_full[id]
|
|
else:
|
|
return None
|
|
|
|
def get_stage_name(self, id: int):
|
|
if id in self.stages:
|
|
return self.stages[id]
|
|
else:
|
|
return None
|
|
|
|
def republish_stages(self):
|
|
self.bus.publish_deferred("pipeline_stages", self.stages)
|
|
|
|
def remove_stage(self, id: int):
|
|
del self.stages[id]
|
|
del self.stagedata[id]
|
|
self.republish_stages()
|