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

import sys

from PyQt5.QtCore import pyqtSignal
from PyQt5.QtGui import QColor, QImage, QPalette, QPixmap
from PyQt5.QtWidgets import QHeaderView, QProgressBar, QTableWidgetItem
from ui.helpers import calc_foreground_color
from ui.test_test import Test_Test


class Test_Vision(Test_Test):
    request_frame = pyqtSignal()

    def __init__(self, components=None, recipe=None, step=None, run_once=False, enable_override=True):
        super().__init__(components=components, recipe=recipe, step=step, run_once=run_once, enable_override=enable_override)
        if "--sim-camera" not in sys.argv:
            self.ok_counter_limit = 2
        else:
            self.ok_counter_limit = 1
        # DETECTIONS TABLE
        self.results_table_headers = {
            "class": "RILEVATO",
            "expected": "ATTESO",
            "score": "PUNTEGGIO",
        }
        self.results_tw.setColumnCount(len(self.results_table_headers))
        self.results_tw.setHorizontalHeaderLabels(self.results_table_headers.values())
        self.results_tw.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.results_tw.verticalHeader().setSectionResizeMode(QHeaderView.Stretch)

    def start(self, recipe=None, step=None):
        super().start(recipe=recipe, step=step)
        if self.run_once is True and self.done is True and self.done_ok is True:
            return
        # 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)
        if "neo_pixels" in self.components and self.components["neo_pixels"].ready:
            self.components["neo_pixels"].set_all_pixel_color("#ffffff")
        self.components["vision"].set_recipe(self.step.spec.get("recipe", None))
        self.components["vision"].resume()
        if "uvc_camera" in self.components:
            self.components["uvc_camera"].resume()
        elif "galaxy_camera" in self.components:
            self.components["galaxy_camera"].resume()
        self.request_frame.emit()  # request first frame
        # DETECTIONS TABLE
        # hide previous results
        self.results_tw.setRowCount(0)

    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 reset(self):
    #     super().reset()

    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)

    @staticmethod
    def tr(text):
        translations = {
            "no_detection": "non_rilevato",
            "_": " ",
            "big": "grande",
            "black": "nero",
            "blue": "blu",
            "brown": "marrone",
            "empty": "vuoto",
            "green": "verde",
            "ko": "ko",
            "ok": "ok",
            "orange": "arancione",
            "red": "rosso",
            "small": "piccolo",
            "white": "bianco",
            "yellow": "giallo",
        }
        for s, t in translations.items():
            text = text.replace(s, t)
        return text

    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)
        results = data.get("results", {}).get("results", 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):
                if "neo_pixels" in self.components and self.components["neo_pixels"].ready:
                    self.components["neo_pixels"].set_all_pixel_color("#ffff00")
            elif self.ok_counter >= self.ok_counter_limit:
                if "neo_pixels" in self.components and self.components["neo_pixels"].ready:
                    self.components["neo_pixels"].set_all_pixel_color("#00ff00")
            else:
                if "neo_pixels" in self.components and self.components["neo_pixels"].ready:
                    self.components["neo_pixels"].set_all_pixel_color("#ff0000")
        else:
            if "neo_pixels" in self.components and self.components["neo_pixels"].ready:
                self.components["neo_pixels"].set_all_pixel_color("#ffffff")
        # DETECTIONS TABLE
        if results is not None:
            self.results_tw.setRowCount(len(results))
            self.results_tw.setVerticalHeaderLabels(list(results))
            for r, [zone_name, result] in enumerate(results.items()):
                for c, header in enumerate(self.results_table_headers):
                    if header == "class":
                        i = self.results_tw.item(r, c)
                        if i is None:
                            i = QTableWidgetItem()
                            self.results_tw.setItem(r, c, i)
                        detection = result.get("detection", None)
                        if detection is not None:
                            text = self.tr(detection["class"]["name"])
                            color = QColor(detection["class"]["color"])
                        else:
                            text = ""
                            color = QColor("#ffffff")
                        i.setText(text)
                        i.setBackground(color)
                        i.setForeground(calc_foreground_color(color))
                    elif header == "expected":
                        i = self.results_tw.item(r, c)
                        if i is None:
                            i = QTableWidgetItem()
                            self.results_tw.setItem(r, c, i)
                        expected = result.get("expected", None)
                        if expected is not None:
                            text = self.tr(expected["name"])
                            color = QColor(expected["color"])
                        else:
                            text = ""
                            color = QColor("#ffffff")
                        i.setText(text)
                        i.setBackground(color)
                        i.setForeground(calc_foreground_color(color))
                    elif header == "score":
                        w = self.results_tw.cellWidget(r, c)
                        if w is None:
                            w = QProgressBar()
                            w.setRange(0, 100)
                            self.results_tw.setCellWidget(r, c, w)
                        score = result.get("detection", {}).get("score", None)
                        ok = result.get("ok", None)
                        if score is not None:
                            value = int(score * 100)
                        else:
                            value = 0
                        w.setValue(value)
                        if ok is not None:
                            if ok:
                                if score is not None:
                                    color = QColor(int(255 - max(score - 0.5, 0) / 0.5 * 255), 255, 0)
                                else:
                                    color = QColor("#00ff00")
                            else:
                                color = QColor("#ff0000")
                        else:
                            color = QColor("#0000ff")
                        p = QPalette()
                        p.setColor(QPalette.Highlight, color)
                        w.setPalette(p)
                    else:
                        raise NotImplementedError()
        else:
            self.results_tw.setRowCount(0)
        super().visualize(data, img=img)

    def save_last(self):
        if self.last is None:
            return
        self.components["vision_saver"].save(
            save_time=self.last.get("time", None),
            frame=self.last.get("frame", None),
            vision=self.last.get("detections", None),
            suffix="manual_save",
            resize=[256, 256],
        )

    def emit_ok(self):
        super().emit_ok()
        if "neo_pixels" in self.components and self.components["neo_pixels"].ready:
            self.components["neo_pixels"].set_all_pixel_color("#000000")

    def emit_ko(self):
        super().emit_ko()
        if "neo_pixels" in self.components and self.components["neo_pixels"].ready:
            self.components["neo_pixels"].set_all_pixel_color("#000000")