import sys from PyQt5.QtCore import pyqtSignal from PyQt5.QtGui import QImage, QPixmap from ui.test_test import Test_Test class Test_Vision(Test_Test): request_frame = pyqtSignal() def __init__(self, components=None, recipe=None, step=None): if "--sim-camera" not in sys.argv: self.ok_counter_limit = 2 else: self.ok_counter_limit = 1 super().__init__(components=components, recipe=recipe, step=step) def start(self, recipe=None, step=None): super().start(recipe=recipe, step=step) # setup camera-vision loop self.components["galaxy_camera"].set_period(period=None) # only get frame on request self.components["galaxy_camera"].add_sources({"test_vision": self.request_frame}) self.request_frame_connection = self.components["vision"].out.connect(self.request_frame) # request new frame as soon as vision finishes self.get_connection = self.components["vision"].out.connect(self.get) self.components["neo_pixels"].set_all_pixel_color("#ffffff") self.components["vision"].set_recipe(self.step.spec.get("recipe", None)) self.components["vision"].resume() self.components["galaxy_camera"].resume() self.request_frame.emit() # request first frame def stop(self): # disable camera-vision loop self.components["galaxy_camera"].pause() self.components["vision"].pause() self.disconnect(self.get_connection) self.disconnect(self.request_frame_connection) self.components["galaxy_camera"].remove_sources(["test_vision", ]) super().stop() def get(self, data=None, override=False): if self.done: # avoid proccessing if completed return if data is None or data[-1] is None: super().get(None, override=override) return data = data[-1] super().get([{ "time": data.get("time", None), "frame": data["vision"].get("frame", None), "detections": data["vision"].get("detections", None), "results": data["vision"].get("results", None), "render": data["vision"].get("render", None), }], override=override) def visualize(self, data=None): if data is None: data = {} overridden = data.get("overridden", False) frame = data.get("frame", None) render = data.get("render", None) if overridden: img = self.status_imgs_full["warning"] elif render is not None: img = render elif frame is not None: img = QPixmap.fromImage(QImage( frame.data, frame.shape[1], # width frame.shape[0], # height frame.shape[2] * frame.shape[1], # width * channels QImage.Format_RGB888 )) else: img = self.status_imgs_full[None] if self.done: if data.get("overridden", False): self.components["neo_pixels"].set_all_pixel_color("#ffff00") elif self.ok_counter >= self.ok_counter_limit: self.components["neo_pixels"].set_all_pixel_color("#00ff00") else: self.components["neo_pixels"].set_all_pixel_color("#ff0000") else: self.components["neo_pixels"].set_all_pixel_color("#ffffff") super().visualize(data, img=img) def save_last(self): if self.last is None: return self.components["vision_saver"].save(self.last["time"], self.last["frame"]) def emit_ok(self): super().emit_ok() self.components["neo_pixels"].set_all_pixel_color("#000000") def emit_ko(self): super().emit_ko() self.components["neo_pixels"].set_all_pixel_color("#000000")