235 lines
10 KiB
Python
235 lines
10 KiB
Python
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, pieces=None, run_once=False, enable_override=True):
|
|
super().__init__(components=components, recipe=recipe, step=step, pieces=pieces, 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, pieces=None):
|
|
show = super().start(recipe=recipe, step=step, pieces=pieces)
|
|
if show is False:
|
|
return show
|
|
# setup camera-vision loop
|
|
if "uvc_camera" in self.components:
|
|
self.components["uvc_camera"].set_period(period=None) # only get frame on request
|
|
self.components["uvc_camera"].add_sources({"test_vision": self.request_frame})
|
|
elif "galaxy_camera" in self.components:
|
|
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)
|
|
return show
|
|
|
|
def stop(self):
|
|
# disable camera-vision loop
|
|
if "uvc_camera" in self.components:
|
|
self.components["uvc_camera"].pause()
|
|
elif "galaxy_camera" in self.components:
|
|
self.components["galaxy_camera"].pause()
|
|
self.components["vision"].pause()
|
|
self.disconnect(self.get_connection)
|
|
self.disconnect(self.request_frame_connection)
|
|
if "uvc_camera" in self.components:
|
|
self.components["uvc_camera"].remove_sources(["test_vision", ])
|
|
elif "galaxy_camera" in self.components:
|
|
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")
|