st-ten-1/src/ui/test_vision/test_vision.py

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["vision"].set_recipe(self.step.spec.get("recipe", "1.ini"))
        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.ok_timer.isActive():
            # avoid proccessing if vision was 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]
        super().visualize(data, img=img)

    def save_last(self):
        if self.last is None:
            return

    def emit_ok(self):
        self.last.pop("frame", None)
        self.last.pop("render", None)
        self.ok.emit(self.last)