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

import sys

from lib.helpers import timing
from PyQt5.QtCore import Qt, QTimer, pyqtSignal
from PyQt5.QtGui import QColor, QImage, QPalette, QPixmap
from ui import Dialog
from ui.test_admin_permission import Test_Admin_Permission
from ui.widget import Widget


class Test_Vision(Widget):
    ok = pyqtSignal(dict)
    request_frame = pyqtSignal()

    def __init__(self, components, recipe):
        super().__init__()
        self.components = components
        self.recipe = recipe
        # setup vision variables
        self.vision = {}
        self.vision_ok_counter = 0
        if "--sim-camera" not in sys.argv:
            self.vision_ok_counter_limit = 2
        else:
            self.vision_ok_counter_limit = 1
        # setup vision controls
        self.ok_timer = QTimer()
        self.ok_timer.setSingleShot(True)
        self.ok_timer.setInterval(2000)
        self.ok_timer.timeout.connect(self.emit_ok)
        # setup save frame button
        self.last_vision = None
        self.save_frame_b.setEnabled(False)
        self.save_frame_b.clicked.connect(self.save_last_vision)
        # setup vision override
        self.admin_challenged = False
        self.vision_overridden = False
        self.override_b.clicked.connect(self.override_vision)
        self.override_b.setEnabled(True)
        # setup vision staus gui
        self.status_imgs_full = {
            True: QPixmap("src/ui/imgs/success.png"),
            "": QPixmap("src/ui/imgs/neo.ico"),
            "warning": QPixmap("src/ui/imgs/warning.png"),
            False: QPixmap("src/ui/imgs/fail.png"),
            None: QPixmap("src/ui/imgs/wait.png"),
        }
        self.status_imgs_small = {k: i.scaled(32, 32, Qt.KeepAspectRatio, Qt.SmoothTransformation) for k, i in self.status_imgs_full.items()}
        self.status_palettes = {
            True: QPalette(),
            "": QPalette(),
            "warning": QPalette(),
            False: QPalette(),
            None: QPalette(),
        }
        self.status_palettes[True].setColor(QPalette.Base, Qt.green)
        self.status_palettes[False].setColor(QPalette.Base, Qt.red)
        self.status_palettes["warning"].setColor(QPalette.Base, QColor(255, 165, 0))
        self.status_palettes[""].setColor(QPalette.Base, QColor(255, 255, 0))
        self.visualize_vision()

    def start(self):
        self.visualize_vision()
        # TESTING
        if "--test" in sys.argv:
            self.override_b.click()
        # /TESTING
        # 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.process_vision_connection = self.components["vision"].out.connect(self.process_vision)
        self.components["vision"].resume()
        self.components["galaxy_camera"].resume()
        # start test
        self.start_time = timing()
        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.process_vision_connection)
        self.disconnect(self.request_frame_connection)
        self.components["galaxy_camera"].remove_sources(["test_vision", ])

    def process_vision(self, data=None, override=False):
        if self.ok_timer.isActive():
            # avoid proccessing if vision was completed
            return
        if data is None:
            data = [{"vision": {}}]
        data = data[-1]
        time = data.get("time", None)
        data = data["vision"]
        frame = data.get("frame", None)
        detections = data.get("detections", None)
        results = data.get("results", None)
        render = data.get("render", None)
        if not override:
            result_ok = data.get("vision", {}).get("ok", False)
        else:
            result_ok = True
        self.last_vision = {
            "time": time,
            "frame": frame,
            "detections": detections,
            "results": results,
            "render": render,
            "overridden": override,
            "vision_duration": timing() - self.start_time,
            "ok": result_ok,
        }
        if not override:
            if result_ok is True:
                self.vision_ok_counter += 1
            else:
                self.vision_ok_counter = 0
        else:
            self.vision_ok_counter = self.vision_ok_counter_limit
        # check if completed
        if self.vision_ok_counter >= self.vision_ok_counter_limit:
            self.stop()
            self.ok_timer.start()
        self.visualize_vision(
            time=time,
            frame=frame,
            detections=detections,
            results=results,
            render=render,
            overridden=override,
        )

    def visualize_vision(self, time=None, frame=None, detections=None, results=None, render=None, overridden=False):
        self.save_frame_b.setEnabled(self.last_vision is not None)
        if overridden:
            self.state_l.setPixmap(self.status_imgs_small["warning"])
        elif results is None:
            self.state_l.setPixmap(self.status_imgs_small[None])
        elif results.get("ok", False) is True:
            self.state_l.setPixmap(self.status_imgs_small[True])
        else:
            self.state_l.setPixmap(self.status_imgs_small[False])
        self.ok_counter_pb.setMaximum(self.vision_ok_counter_limit)
        self.ok_counter_pb.setValue(min(self.vision_ok_counter, self.vision_ok_counter_limit))
        if self.vision_ok_counter >= self.vision_ok_counter_limit:
            self.ok_counter_pb.setPalette(self.status_palettes[True])
        else:
            self.ok_counter_pb.setPalette(self.status_palettes[False])
        if overridden:
            self.img = self.status_imgs_full["warning"]
        elif render is not None:
            self.img = render
        elif frame is not None:
            self.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:
            self.img = self.status_imgs_full[None]
        self.resizeEvent()

    def save_last_vision(self):
        if self.last_vision is None:
            return

    def resizeEvent(self, event=None):
        self.frame_l.setPixmap(self.img.scaled(self.frame_l.width(), self.frame_l.height(), Qt.KeepAspectRatio, Qt.SmoothTransformation))

    def challenge_admin(self, info):
        if not self.admin_challenged:
            d = Dialog()
            d.setCentralWidget(Test_Admin_Permission(info))
            d.setModal(True)
            d.show()
            r = d.exec()
            if r == d.Accepted:
                self.admin_challenged = True
            elif r == d.Rejected:
                self.admin_challenged ^= False
            else:
                raise AssertionError("Bad admin challenge dialog return code")
        return self.admin_challenged

    def override_vision(self):
        if self.challenge_admin("Si sta tentando di bypassare il test di visione"):
            self.process_vision(override=True)

    def emit_ok(self):
        self.ok.emit(self.last_vision)