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: 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 # 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: self.components["neo_pixels"].set_all_pixel_color("#ffff00") elif self.ok_counter >= self.ok_counter_limit: if "neo_pixels" in self.components: self.components["neo_pixels"].set_all_pixel_color("#00ff00") else: if "neo_pixels" in self.components: self.components["neo_pixels"].set_all_pixel_color("#ff0000") else: if "neo_pixels" in self.components: 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: self.components["neo_pixels"].set_all_pixel_color("#000000") def emit_ko(self): super().emit_ko() if "neo_pixels" in self.components: self.components["neo_pixels"].set_all_pixel_color("#000000")