st-ten-1/src/components/vision.py

792 lines
35 KiB
Python
Raw Normal View History

2022-07-18 13:03:54 +00:00
import logging
2022-06-21 12:10:52 +00:00
import os
import sys
import traceback
from configparser import ConfigParser
from pathlib import Path
2022-06-27 09:03:26 +00:00
import cv2
2022-06-21 12:10:52 +00:00
import numpy as np
2022-06-28 10:31:27 +00:00
from PyQt5.QtCore import (QFileSystemWatcher, QLineF, QMutex, QPointF, QRectF,
2022-07-06 19:37:08 +00:00
Qt, QThread, pyqtSignal)
2022-06-28 10:31:27 +00:00
from PyQt5.QtGui import (QBrush, QColor, QFont, QImage, QPainter, QPainterPath,
QPen, QPixmap)
2022-07-20 17:29:11 +00:00
from PyQt5.QtWidgets import QApplication
2022-06-21 12:10:52 +00:00
from .component import Component
2022-06-22 15:18:29 +00:00
from .consumer import Consumer
2022-06-21 12:10:52 +00:00
2022-07-18 13:52:23 +00:00
if "--no-gpu" in sys.argv:
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import tensorflow as tf
from lib.helpers.object_detection.utils import label_map_util
2022-06-21 12:10:52 +00:00
if "--no-edgetpu" not in sys.argv:
2022-07-04 10:36:51 +00:00
try:
2022-06-21 12:10:52 +00:00
from pycoral.utils.edgetpu import make_interpreter
2022-07-06 19:37:08 +00:00
except (ImportError, ModuleNotFoundError):
2022-07-18 12:45:24 +00:00
logging.exception(traceback.format_exc())
2022-07-04 10:36:51 +00:00
def make_interpreter(*args, **kwargs):
2022-07-06 19:37:08 +00:00
raise ValueError("the 'pycoral' module is not available")
2022-06-21 12:10:52 +00:00
else:
def make_interpreter(*args, **kwargs):
raise ValueError("\"--no-edgetpu\" in sys.argv")
if "--no-tflite" not in sys.argv:
2022-07-04 10:36:51 +00:00
try:
2022-07-18 13:03:54 +00:00
Interpreter = tf.lite.Interpreter
# from tflite_runtime.interpreter import Interpreter
2022-07-06 19:37:08 +00:00
except (ImportError, ModuleNotFoundError):
2022-07-18 12:45:24 +00:00
logging.exception(traceback.format_exc())
2022-07-04 10:36:51 +00:00
def Interpreter(*args, **kwargs):
2022-07-06 19:37:08 +00:00
raise ValueError("the 'tflite-runtime' module is not available")
2022-06-21 12:10:52 +00:00
else:
def Interpreter(*args, **kwargs):
raise ValueError("\"--no-tflite\" in sys.argv")
2022-08-02 16:15:30 +00:00
if "--fail-vision" not in sys.argv:
vision_override = None
else:
vision_override = False
2022-06-21 12:10:52 +00:00
# # Patch the location of gfile
# tf.gfile = tf.io.gfile
class Vision(Component):
2022-06-28 10:31:27 +00:00
"""everything is expected the have shape with height (y) first then width (x)"""
2022-06-21 12:10:52 +00:00
status_signal = pyqtSignal(dict)
loading_model_signal = pyqtSignal(dict)
2022-06-22 15:18:29 +00:00
def __init__(self, config=None, name=None, period=None, lazy=True, paused=False, threaded=True):
2022-06-21 12:10:52 +00:00
super().__init__(config=config, name=name, period=period, lazy=lazy, paused=paused, threaded=threaded)
self.lock = QMutex()
self.simulate = "--sim-vision" in sys.argv
2022-06-22 15:18:29 +00:00
def start(self):
2022-06-21 12:10:52 +00:00
self.model = None
2022-06-28 10:31:27 +00:00
# VISION THREAD
self.vision_consumer = Consumer(work=self.vision_consumer_work, work_fifo=True, drop_fifo=True, work_maxlen=1, name="vision_consumer", paused=False)
self.vision_consumer_thread = QThread()
self.vision_consumer_thread.setTerminationEnabled(True)
self.vision_consumer.moveToThread(self.vision_consumer_thread)
self.vision_consumer_thread.started.connect(self.vision_consumer.start)
self.vision_consumer_thread.start()
2022-07-20 17:29:11 +00:00
if "--debugger-workaround" in sys.argv:
QApplication.processEvents()
QThread.msleep(1000)
QApplication.processEvents()
2022-06-28 10:31:27 +00:00
self.vision_consumer.wait_ready()
# RENDER THREAD
self.render_consumer = Consumer(work=self.render_consumer_work, work_fifo=True, drop_fifo=True, work_maxlen=1, name="render_consumer", paused=False)
self.render_consumer_thread = QThread()
self.render_consumer_thread.setTerminationEnabled(True)
self.render_consumer.moveToThread(self.render_consumer_thread)
self.render_consumer_thread.started.connect(self.render_consumer.start)
self.render_consumer_thread.start()
2022-07-20 17:29:11 +00:00
if "--debugger-workaround" in sys.argv:
QApplication.processEvents()
QThread.msleep(1000)
QApplication.processEvents()
2022-06-28 10:31:27 +00:00
self.render_consumer.wait_ready()
# CONNECT CONSUMERS
self.vision_consumer.out.connect(self.process_vision_consumed)
self.render_consumer.out.connect(self.process_render_consumed)
2022-06-22 15:18:29 +00:00
super().start()
2022-06-21 12:10:52 +00:00
def config_changed(self):
# OBJECT DETECTION
2022-08-02 16:15:30 +00:00
self.detection_threshold = float(self.config[self.name].get("detection_threshold", 0.5))
2022-06-21 12:10:52 +00:00
# recipe
self.zones = None
self.labels = None
# LOAD RECIPE
self.recipes_dir = Path(self.config[self.name].get("recipes_dir", "./config/vision/recipes"))
self.set_recipe(None)
self.recipe_watcher = QFileSystemWatcher([])
self.recipe_watcher.fileChanged.connect(self._set_recipe)
# LOAD MODEL
self.models_dir = Path(self.config[self.name].get("models_dir", "./data/neural_networks"))
self.allowed_modes = dict.fromkeys([
"edgetpu",
2022-07-06 19:37:08 +00:00
"tflite",
2022-06-21 12:10:52 +00:00
"normal",
])
if "--no-edgetpu" in sys.argv:
self.allowed_modes.pop("edgetpu", None)
if "--no-tflite" in sys.argv:
self.allowed_modes.pop("edgetpu", None)
2022-07-06 19:37:08 +00:00
self.allowed_modes.pop("tflite", None)
2022-06-21 12:10:52 +00:00
self.load_model(self.config[self.name].get("neural_network", None))
# LOAD LABELS
label_map = label_map_util.load_labelmap("./config/vision/labels/labels.pbtxt")
self.num_classes = len(label_map.item)
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=self.num_classes, use_display_name=True)
self.category_index = label_map_util.create_category_index(categories)
2022-08-02 16:15:30 +00:00
for k in self.category_index:
self.category_index[k]["color"] = self.category_index[k]["color"].replace("0x", "#")
2022-06-21 12:10:52 +00:00
self.classes_map = {c["name"]: k for k, c in self.category_index.items()}
2022-08-02 16:15:30 +00:00
self.zone_detection_filter_mode = self.config[self.name].get("zone_detection_filter_mode", "box_inside")
self.zone_detection_preference_mode = self.config[self.name].get("zone_detection_preference_mode", "score")
2022-06-28 10:31:27 +00:00
2022-08-02 16:15:30 +00:00
@staticmethod
def get_center(rect):
2022-06-28 10:31:27 +00:00
return [(rect[0] + rect[2]) / 2, (rect[1] + rect[3]) / 2]
2022-08-02 16:15:30 +00:00
@staticmethod
def get_size(rect):
2022-06-28 10:31:27 +00:00
return [rect[2] - rect[0], rect[3] - rect[1]]
2022-08-02 16:15:30 +00:00
@staticmethod
def get_box(center, size):
2022-06-28 10:31:27 +00:00
return [center[0] - size[0] / 2, center[1] - size[1] / 2, center[0] + size[0] / 2, center[1] + size[1] / 2]
2022-08-02 16:15:30 +00:00
@staticmethod
def get_distance(p1, p2):
2022-06-28 10:31:27 +00:00
return pow((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2, 1 / 2)
2022-06-21 12:10:52 +00:00
def clear_recipe(self):
self.recipe = None
2022-06-28 10:31:27 +00:00
self.markers = {}
2022-06-21 12:10:52 +00:00
self.zones = {}
self.labels = {}
def _set_recipe(self, recipe_path):
recipe_path = str(recipe_path)
self.log.info(f"changing recipe to {recipe_path!r}")
watched = self.recipe_watcher.files()
if recipe_path == "" and len(watched) == 0: # skip bad watcher signals
return
if len(watched) > 0:
self.recipe_watcher.removePaths(watched)
self.recipe_path = recipe_path
try:
if not os.path.isfile(self.recipe_path):
raise AssertionError(f"Recipe file {self.recipe_path!r} could not be found.")
config = ConfigParser(inline_comment_prefixes="#")
read = config.read(self.recipe_path)
if len(read) != 1 or self.recipe_path not in read:
raise AssertionError("Recipe could not be read.")
os.path.splitext(os.path.basename(read[0]))[0]
2022-06-28 10:31:27 +00:00
self.markers = self.parse_markers(config._sections.get("markers", None))
self.zones = self.parse_zones(config._sections.get("zones", None))
self.labels = self.parse_labels(config._sections.get("labels", None))
2022-06-21 12:10:52 +00:00
self.recipe_watcher.addPath(str(self.recipe_path))
except Exception:
self.log.exception(traceback.format_exc())
self.log.exception(f"Error reading {self.recipe_path!r}:")
self.clear_recipe()
self.status_signal.emit({
"recipe": self.recipe,
2022-06-28 10:31:27 +00:00
"markers": self.markers,
2022-06-21 12:10:52 +00:00
"zones": self.zones,
"labels": self.labels,
})
def set_recipe(self, recipe=None):
if recipe is None:
self.clear_recipe()
else:
2022-08-02 16:15:30 +00:00
self.recipe = recipe
2022-06-21 12:10:52 +00:00
self._set_recipe(self.recipes_dir / str(recipe))
2022-06-28 10:31:27 +00:00
def parse_markers(self, config=None):
2022-06-21 12:10:52 +00:00
if config is None:
2022-06-28 10:31:27 +00:00
raise AssertionError(f"Recipe file {self.recipe_path!r} does not contain the 'markers' section.")
2022-06-21 12:10:52 +00:00
config = {}
2022-06-28 10:31:27 +00:00
markers = {}
for marker_name, marker_spec in config.items():
2022-06-21 12:10:52 +00:00
try:
2022-06-28 10:31:27 +00:00
center, size, fill_color, border_color, border_thickness, shape = marker_spec.split(" ")
center = list(reversed(list(map(float, center.split(",")))))
size = list(reversed(list(map(float, size.split(",")))))
2022-06-21 12:10:52 +00:00
if len(size) == 1:
size = [size[0], size[0]]
fill_color = QColor(fill_color.replace("0x", "#"))
border_color = QColor(border_color.replace("0x", "#"))
border_thickness = float(border_thickness)
shape = shape.lower()
2022-06-28 10:31:27 +00:00
markers[marker_name] = {
2022-06-21 12:10:52 +00:00
"border_color": border_color,
"border_thickness": border_thickness,
2022-06-28 10:31:27 +00:00
"center": center,
"fill_color": fill_color,
2022-06-21 12:10:52 +00:00
"shape": shape,
2022-06-28 10:31:27 +00:00
"size": size,
2022-06-21 12:10:52 +00:00
}
except Exception:
self.log.exception(traceback.format_exc())
2022-06-28 10:31:27 +00:00
self.log.exception(f"marker {marker_name!r} in recipe file {self.recipe_path!r} could not be parsed. spec: {marker_spec!r}")
return markers
2022-06-21 12:10:52 +00:00
def parse_zones(self, config=None):
if config is None:
raise AssertionError(f"Recipe file {self.recipe_path!r} does not contain the 'zones' section.")
config = {}
zones = {}
for zone_name, zone_spec in config.items():
zone_name = zone_name.upper()
try:
2022-06-28 10:31:27 +00:00
center, size, d_class = zone_spec.split(" ")
center = list(reversed(list(map(float, center.split(",")))))
if "," in size:
size = list(reversed(list(map(float, size.split(",")))))
shape = "rect"
2022-06-21 12:10:52 +00:00
else:
2022-06-28 10:31:27 +00:00
size = [float(size)] * 2
shape = "ellipse"
d_class = self.category_index[self.classes_map[d_class]]
2022-06-21 12:10:52 +00:00
zones[zone_name] = {
2022-08-02 16:15:30 +00:00
"border_color": QColor(d_class["color"]),
2022-06-28 10:31:27 +00:00
"border_thickness": 25,
"box": self.get_box(center, size),
2022-08-02 16:15:30 +00:00
"convert_negative_placement": False,
2022-06-21 12:10:52 +00:00
"center": center,
2022-06-28 10:31:27 +00:00
"class": d_class,
"fill_color": QColor("#00000000"),
"pen_line": "DashLine",
"shape": shape,
"size": size,
2022-06-21 12:10:52 +00:00
}
except Exception:
self.log.exception(traceback.format_exc())
self.log.exception(f"region {zone_name!r} in recipe file {self.recipe_path!r} could not be parsed. spec: {zone_spec!r}")
return zones
def parse_labels(self, config=None):
if config is None:
config = {}
labels = {}
for label_name, label_spec in config.items():
try:
location, font_size, fill_color, border_color, border_thickness, text = label_spec.split(" ", 5)
2022-06-28 10:31:27 +00:00
location = list(reversed(list(map(float, location.split(",")))))
2022-06-21 12:10:52 +00:00
font_size = float(font_size)
fill_color = QColor(fill_color.replace("0x", "#"))
border_color = QColor(border_color.replace("0x", "#"))
border_thickness = float(border_thickness)
text = text.replace("\\n", "\n").replace("\\t", "\t")
labels[label_name] = {
"border_color": border_color,
"border_thickness": border_thickness,
2022-06-28 10:31:27 +00:00
"fill_color": fill_color,
"font_size": font_size,
"location": location,
"opacity": 1,
"shape": "text",
2022-06-21 12:10:52 +00:00
"text": text,
}
except Exception:
self.log.exception(traceback.format_exc())
self.log.exception(f"label {label_name!r} in recipe file {self.recipe_path!r} could not be parsed. spec: {label_spec!r}")
return labels
def load_model(self, model=None):
2022-06-22 15:18:29 +00:00
# print("VISION CONSUMER", str(int(QThread.currentThreadId())), flush=True)
2022-06-21 12:10:52 +00:00
self.log.info(f"requested neural network: {model!r}")
if model is None or model.lower() in [
"",
"any",
"last",
"latest",
"newest",
"none",
]:
model_name = sorted([d for d in os.listdir(self.models_dir) if os.path.isdir(self.models_dir / d)], reverse=True)[0]
2022-07-18 16:20:16 +00:00
else:
model_name = model
2022-06-21 12:10:52 +00:00
self.log.info(f"loading neural network: {model_name!r}")
self.loading_model_signal.emit({"status": "loading"})
self.lock.lock()
tf_mode = None
# reset tflite variables
interpreter = None
2022-07-18 14:07:10 +00:00
if self.simulate:
tf_mode = "simulation"
interpreter = None
2022-06-21 12:10:52 +00:00
if tf_mode is None and "edgetpu" in self.allowed_modes:
try:
# create tflite edgetpu interpreter
2022-06-27 09:03:26 +00:00
interpreter = make_interpreter(str(self.models_dir / model_name / f"{model_name}_edgetpu.tflite"))
2022-06-21 12:10:52 +00:00
tf_mode = "edgetpu"
except Exception:
self.log.exception(traceback.format_exc())
2022-07-06 19:37:08 +00:00
if tf_mode is None and "tflite" in self.allowed_modes:
2022-06-21 12:10:52 +00:00
try:
# create tflite cpu interpreter
2022-06-27 09:03:26 +00:00
interpreter = Interpreter(str(self.models_dir / model_name / f"{model_name}.tflite"))
2022-07-06 19:37:08 +00:00
tf_mode = "tflite"
2022-06-21 12:10:52 +00:00
except Exception:
self.log.exception(traceback.format_exc())
# reset tensorflow variables
model = None
if tf_mode is None and "normal" in self.allowed_modes:
try:
# create tensorflow model
2022-08-01 11:29:12 +00:00
model = tf.saved_model.load(str(self.models_dir / model_name)).signatures["serving_default"]
2022-06-21 12:10:52 +00:00
tf_mode = "normal"
except Exception:
self.log.exception(traceback.format_exc())
self.lock.unlock()
if tf_mode is None:
raise RuntimeError("failed initialize any neural network model")
self.tf_mode = tf_mode
self.model_name = model_name
if interpreter is not None:
2022-07-06 19:37:08 +00:00
# if there is a new tflite interpreter initialize it
2022-06-27 09:03:26 +00:00
interpreter.allocate_tensors()
interpreter.invoke() # warmup
2022-06-21 12:10:52 +00:00
self.interpreter = interpreter
# if there is a new model to be used, remove previous model if present
if model is not None and self.model is not None:
tf.keras.backend.clear_session()
self.model = model
self.log.info(f"initialized model {self.model!r} with mode {self.tf_mode!r}")
self.loading_model_signal.emit({"status": "done"})
2022-06-28 10:31:27 +00:00
def check_features(self, frame, lock=True):
2022-07-06 19:37:08 +00:00
if self.interpreter is not None and frame.shape != self.interpreter.get_input_details()[0]["shape"][1:3]:
tensor = np.expand_dims(cv2.resize(frame, self.interpreter.get_input_details()[0]["shape"][1:3], interpolation=cv2.INTER_LINEAR), axis=0)
2022-06-28 10:31:27 +00:00
else:
2022-08-01 11:29:12 +00:00
frame_resized = cv2.resize(frame, (256, 256), interpolation=cv2.INTER_LINEAR)
tensor = tf.convert_to_tensor(np.asarray(frame_resized))
tensor = tensor[tf.newaxis, ...]
# tensor = np.expand_dims(frame, axis=0)
2022-06-21 12:10:52 +00:00
# Run inference
if lock:
self.lock.lock()
2022-07-19 09:59:00 +00:00
if self.simulate or self.tf_mode == "simulation":
2022-07-18 14:07:10 +00:00
detections = {
"detection_scores": [[1.0]],
"detection_boxes": [[[0.2, 0.2, 0.8, 0.8]]],
"detection_classes": [[1]],
}
if lock:
self.lock.unlock()
elif self.tf_mode in {"edgetpu", "tflite"}:
2022-07-06 19:37:08 +00:00
i_d = self.interpreter.get_input_details()
# print(i_d)
o_d = self.interpreter.get_output_details()
# print(o_d)
self.interpreter.set_tensor(i_d[0]["index"], tensor)
2022-06-21 12:10:52 +00:00
self.interpreter.invoke()
2022-07-06 19:37:08 +00:00
# PARSE TFLITE DETECTIONS
# signature_list = self.interpreter._get_full_signature_list()
# if signature_list:
# if len(signature_list) > 1:
# raise ValueError("Only support model with one signature.")
# signature = signature_list[next(iter(signature_list))]
# # count = int(self.interpreter.get_tensor(signature["outputs"]["output_0"])[0])
# scores = self.interpreter.get_tensor(signature["outputs"]["output_1"])[0]
# class_ids = self.interpreter.get_tensor(signature["outputs"]["output_2"])[0]
# boxes = self.interpreter.get_tensor(signature["outputs"]["output_3"])[0]
if self.interpreter.get_tensor(o_d[3]["index"]).size == 1:
boxes = self.interpreter.get_tensor(o_d[0]["index"])[0]
class_ids = self.interpreter.get_tensor(o_d[1]["index"])[0]
scores = self.interpreter.get_tensor(o_d[2]["index"])[0]
# count = int(self.interpreter.get_tensor(o_d[3]["index"])[0])
else:
scores = self.interpreter.get_tensor(o_d[0]["index"])[0]
boxes = self.interpreter.get_tensor(o_d[1]["index"])[0]
# count = int(self.interpreter.get_tensor(o_d[2]["index"])[0])
class_ids = self.interpreter.get_tensor(o_d[3]["index"])[0]
2022-06-21 12:10:52 +00:00
if lock:
self.lock.unlock()
detections = {
2022-06-27 09:03:26 +00:00
"detection_scores": [scores],
2022-07-06 19:37:08 +00:00
"detection_boxes": [boxes],
"detection_classes": [map(lambda class_id: class_id + 1, class_ids)],
2022-06-21 12:10:52 +00:00
}
else:
detections = self.model(tensor)
if lock:
self.lock.unlock()
2022-08-01 11:29:12 +00:00
detections = {
"detection_scores": detections["detection_scores"].numpy().tolist(),
"detection_boxes": detections["detection_boxes"].numpy().tolist(),
"detection_classes": detections["detection_classes"],
}
2022-06-27 09:03:26 +00:00
# WARNING: results other than the ones related to tensor[-1] will be discarded
2022-06-21 12:10:52 +00:00
parsed_detections = []
2022-08-01 11:29:12 +00:00
for d_score, d_box, d_class in zip( # , d_mask in zip(
detections["detection_scores"][-1],
2022-06-27 09:03:26 +00:00
detections["detection_boxes"][-1],
detections["detection_classes"][-1],
# detections["detection_masks"][-1],
2022-06-21 12:10:52 +00:00
):
if d_score < self.detection_threshold:
continue
2022-06-27 09:03:26 +00:00
box = list(d_box)
2022-06-28 10:31:27 +00:00
box = [i * s for i, s in zip(box, frame.shape[:2] * 2)] # rescale detection to frame size
2022-06-21 12:10:52 +00:00
detection = {
2022-08-01 11:29:12 +00:00
"score": d_score,
2022-06-27 09:03:26 +00:00
"box": box,
2022-08-01 11:29:12 +00:00
"class": self.category_index[int(d_class)],
# "mask": d_mask,
2022-06-28 10:31:27 +00:00
"center": self.get_center(box),
"size": self.get_size(box),
2022-06-21 12:10:52 +00:00
}
parsed_detections.append(detection)
2022-06-28 10:31:27 +00:00
return parsed_detections
2022-06-21 12:10:52 +00:00
2022-06-28 10:31:27 +00:00
def detections_to_items(self, detections):
# DRAW DETECTIONS
if detections is not None and len(detections):
style = {
"border_thickness": 25,
"fill_color": QColor("#00000000"),
2022-08-02 16:15:30 +00:00
"shape": "rect",
"convert_negative_placement": False,
2022-06-28 10:31:27 +00:00
}
items = {}
for item_name, item in enumerate(detections):
items[str(item_name)] = {
2022-08-02 16:15:30 +00:00
**item,
2022-06-28 10:31:27 +00:00
**style,
2022-08-02 16:15:30 +00:00
"border_color": QColor(item["class"]["color"]),
2022-06-28 10:31:27 +00:00
}
return items
else:
return {}
def process_detections(self, detections):
2022-08-02 16:15:30 +00:00
if self.zones is None or not len(self.zones):
return None
2022-06-28 10:31:27 +00:00
# MATCH DETECTIONS WITH RECIPE
results = dict.fromkeys(self.zones)
for detection in detections:
# find closest zone center to the detection
# filtering out those that do not contain the detection
min_distance = sys.maxsize
closest_zone = None
for zone_name, zone in self.zones.items():
2022-07-20 17:29:11 +00:00
distance = self.get_distance(detection["center"], zone["center"])
2022-06-28 10:31:27 +00:00
if zone["shape"] == "rect":
if self.zone_detection_filter_mode == "center_inside":
outside_zone = any([
detection["center"][0] < zone["box"][0],
detection["center"][0] > zone["box"][2],
detection["center"][1] < zone["box"][1],
detection["center"][1] > zone["box"][3],
])
elif self.zone_detection_filter_mode == "box_inside":
outside_zone = any([
detection["box"][0] < zone["box"][0],
detection["box"][2] > zone["box"][2],
detection["box"][1] < zone["box"][1],
detection["box"][3] > zone["box"][3],
])
elif self.zone_detection_filter_mode == "box_touches":
outside_zone = any([
detection["box"][2] < zone["box"][0],
detection["box"][0] > zone["box"][2],
detection["box"][3] < zone["box"][1],
detection["box"][1] > zone["box"][3],
])
else:
raise NotImplementedError(f"invalid zone_detection_filter_mode: {self.zone_detection_filter_mode!r}")
elif zone["shape"] == "ellipse": # it's a circle
if self.zone_detection_filter_mode == "center_inside":
outside_zone = distance > zone["size"][0] / 2
elif self.zone_detection_filter_mode == "box_inside":
outside_zone = distance + max(detection["size"]) / 2 > zone["size"][0] / 2
elif self.zone_detection_filter_mode == "box_touches":
outside_zone = distance - max(detection["size"]) / 2 > zone["size"][0] / 2
2022-07-20 17:29:11 +00:00
else:
raise NotImplementedError(f"invalid zone_detection_filter_mode: {self.zone_detection_filter_mode!r}")
2022-06-28 10:31:27 +00:00
else:
raise NotImplementedError(f"invalid zone shape: {zone['shape']!r}")
if not outside_zone and distance < min_distance:
min_distance = distance
closest_zone = zone_name
if closest_zone is not None:
# if closest zone already has a matching detection
# replace it only if the current one is preferred
if results[closest_zone] is not None:
if self.zone_detection_preference_mode == "distance":
if min_distance >= results[closest_zone]["zone_distance"]:
continue
elif self.zone_detection_preference_mode == "score":
if detection["score"] <= results[closest_zone]["score"]:
continue
else:
raise NotImplementedError(f"invalid zone_detection_preference_mode: {self.zone_detection_preference_mode!r}")
results[closest_zone] = detection.copy()
results[closest_zone]["zone_distance"] = min_distance
# check detections against recipe
checked = {}
for zone_name, detection in results.items():
if detection is None:
# dummy empty detection if nothing detected
detection = {
"box": [0, 0, 0, 0],
"center": [0, 0],
"class": {
"id": None,
"name": "no_detection",
"color": "rgb(0,0,0)",
},
# "mask": [],
"score": 1,
"size": [0, 0],
}
if zone_name not in self.zones or self.zones[zone_name]["class"]["id"] in {"no_detection", "none", }:
expected_class = {
"id": None,
"name": "no_detection",
"color": "rgb(0,0,0)",
}
else:
expected_class = self.zones[zone_name]["class"]
checked[zone_name] = {
"ok": detection is not None and detection["class"]["id"] == expected_class["id"],
"expected": expected_class,
"detection": detection,
}
2022-08-02 16:15:30 +00:00
global vision_override
if vision_override is None:
ok = all(map(lambda detection: detection["ok"] is True, checked.values()))
else:
ok = vision_override
2022-06-28 10:31:27 +00:00
return {
2022-08-02 16:15:30 +00:00
"ok": ok,
2022-06-28 10:31:27 +00:00
"results": checked,
}
def results_to_items(self, results, ):
# DRAW ZONES RESULTS
if self.zones is not None and len(self.zones) and results is not None and len(results):
style = {
2022-08-02 16:15:30 +00:00
"pen_line": "SolidLine",
2022-06-28 10:31:27 +00:00
"border_thickness": 50,
"fill_color": QColor("#00000000"),
}
items = {
"_global_result": {
"box": [
style["border_thickness"],
style["border_thickness"],
-style["border_thickness"] * 2,
-style["border_thickness"] * 2,
],
"shape": "rect",
**style,
"border_color": Qt.green if results["ok"] else Qt.red,
}
}
for item_name, item in results["results"].items():
zone = self.zones[item_name]
items[str(item_name)] = {
2022-08-02 16:15:30 +00:00
**zone,
2022-06-28 10:31:27 +00:00
**style,
"border_color": Qt.green if item["ok"] else Qt.red,
}
return items
else:
return {}
2022-08-02 16:15:30 +00:00
@staticmethod
def convert_negative_placement(p, painter):
for k in p:
if p[k] < 0:
if k.startswith("x") or k.startswith("w"):
p[k] += painter.device().width()
elif k.startswith("y") or k.startswith("h"):
p[k] += painter.device().height()
else:
raise AssertionError("could not detect variable direction")
@staticmethod
def apply_placement_offset(p, offset):
for k in p:
if k.startswith("x"):
p[k] += offset["x"]
elif k.startswith("y"):
p[k] += offset["y"]
else:
raise AssertionError("could not detect variable direction")
2022-06-28 10:31:27 +00:00
def render_items(self, items, offset=None, qimage=None, painter=None):
if offset is None:
2022-08-02 16:15:30 +00:00
offset = {"x": 0, "y": 0}
2022-06-28 10:31:27 +00:00
if painter is None:
if qimage is None:
raise AssertionError("one of 'qimage' or 'painter' parameter must not be None")
painter = QPainter()
painter.begin(qimage)
for item_name, item in items.items():
try:
2022-08-02 16:15:30 +00:00
convert_negative_placement = item.get("convert_negative_placement", True)
2022-06-28 10:31:27 +00:00
if "box" in item:
2022-08-02 16:15:30 +00:00
p = {"x1": item["box"][1], "y1": item["box"][0], "x2": item["box"][3], "y2": item["box"][2], }
if convert_negative_placement:
Vision.convert_negative_placement(p, painter)
Vision.apply_placement_offset(p, offset)
p["w"], p["h"] = p["x2"] - p["x1"], p["y2"] - p["y1"]
p["xc"], p["yc"] = Vision.get_center([p["x1"], p["y1"], p["x2"], p["y2"]])
2022-06-28 10:31:27 +00:00
elif "location" in item:
2022-08-02 16:15:30 +00:00
p = {"x1": item["location"][1], "y1": item["location"][0], }
if convert_negative_placement:
Vision.convert_negative_placement(p, painter)
Vision.apply_placement_offset(p, offset)
2022-06-28 10:31:27 +00:00
if "size" in item:
2022-08-02 16:15:30 +00:00
p["w"], p["h"] = item["size"][1], item["size"][0]
p["x2"], p["y2"] = p["x1"] + p["w"], p["y1"] + p["h"]
p["xc"], p["yc"] = Vision.get_center([p["x1"], p["y1"], p["x2"], p["y2"]])
2022-06-28 10:31:27 +00:00
else:
2022-08-02 16:15:30 +00:00
p["w"], p["h"] = 0, 0
p["x2"], p["y2"] = p["x1"], p["y1"]
p["xc"], p["yc"] = p["x1"], p["y1"]
2022-06-28 10:31:27 +00:00
elif "center" in item and "size" in item:
2022-08-02 16:15:30 +00:00
p = {"xc": item["center"][1], "yc": item["center"][0], }
if convert_negative_placement:
Vision.convert_negative_placement(p, painter)
Vision.apply_placement_offset(p, offset)
p["w"], p["h"] = item["size"][1], item["size"][0]
p["x1"], p["y1"], p["x2"], p["y2"] = Vision.get_box([p["xc"], p["yc"]], [p["w"], p["h"]])
2022-06-28 10:31:27 +00:00
else:
raise AssertionError("item has no valid positioning information")
painter.setOpacity(item.get("opacity", 0.5))
painter.setBrush(QBrush(item.get("fill_color", QColor("#ffffff")), getattr(Qt, item.get("brush_pattern", "SolidPattern"))))
painter.setPen(QPen(
item.get("border_color", QColor("#000000")),
item.get("border_thickness", 1),
getattr(Qt, item.get("pen_line", "SolidLine")),
getattr(Qt, item.get("pen_cap", "SquareCap")),
getattr(Qt, item.get("pen_join", "MiterJoin")),
))
if item["shape"] == "ellipse":
2022-08-02 16:15:30 +00:00
painter.drawEllipse(QPointF(p["xc"], p["yc"]), p["w"] / 2, p["h"] / 2)
2022-06-28 10:31:27 +00:00
elif item["shape"] == "cross":
2022-08-02 16:15:30 +00:00
painter.drawLine(QLineF(p["xc"], p["y1"], p["xc"], p["y2"]))
painter.drawLine(QLineF(p["x1"], p["yc"], p["x2"], p["yc"]))
2022-06-28 10:31:27 +00:00
elif item["shape"] == "line":
2022-08-02 16:15:30 +00:00
painter.drawLine(QLineF(p["xc"], p["yc"], p["x2"], p["y2"]))
2022-06-28 10:31:27 +00:00
elif item["shape"] == "rect":
2022-08-02 16:15:30 +00:00
painter.drawRect(QRectF(QPointF(p["x1"], p["y1"]), QPointF(p["x2"], p["y2"])))
2022-06-28 10:31:27 +00:00
elif item["shape"] == "text":
old_render_hints = painter.renderHints()
painter.setRenderHints(QPainter.Antialiasing | QPainter.TextAntialiasing, True)
font = QFont()
font.setStyleHint(QFont.SansSerif, QFont.PreferDefault | QFont.PreferAntialias)
font.setBold(item.get("font_bold", False))
font.setItalic(item.get("font_italic", False))
font.setKerning(item.get("font_kerning", True))
font.setLetterSpacing(QFont.AbsoluteSpacing, item.get("font_letter_spacing", 0))
font.setWordSpacing(item.get("font_word_spacing", 0))
2022-07-18 13:17:49 +00:00
font.setPixelSize(round(item.get("font_size", 25)))
2022-06-28 10:31:27 +00:00
path = QPainterPath()
2022-08-02 16:15:30 +00:00
path.addText(p["x1"], p["y1"], font, item["text"])
2022-06-28 10:31:27 +00:00
painter.drawPath(path)
painter.setRenderHints(old_render_hints, True)
else:
raise NotImplementedError(f"item {item_name!r} has an invalid shape: {item['shape']!r}")
except Exception:
self.log.exception("".join(traceback.format_stack()))
self.log.exception(traceback.format_exc())
self.log.error(f"item {item_name!r} could not be drawn.")
def render(self, frame, detections=None, vision_results=None, mask=True, offset=None):
if mask is True:
qframe = QImage(
frame.shape[1], # width
frame.shape[0], # height
QImage.Format_RGBA8888
)
else:
aframe = cv2.cvtColor(frame, cv2.COLOR_RGB2RGBA)
qframe = QImage(
aframe.data,
aframe.shape[1], # width
aframe.shape[0], # height
aframe.shape[2] * frame.shape[1], # width * channels
QImage.Format_RGBA8888
)
painter = QPainter()
painter.begin(qframe)
for items in [
self.markers,
self.zones,
self.labels,
self.detections_to_items(detections),
self.results_to_items(vision_results)
]:
if items is not None:
try:
self.render_items(
items,
offset=offset,
painter=painter,
)
except Exception:
self.log.exception(traceback.format_exc())
painter.end()
return qframe
def _get(self, data=None):
2022-06-21 12:10:52 +00:00
# print("VISION", str(int(QThread.currentThreadId())), flush=True)
2022-06-28 10:31:27 +00:00
if data is None:
return
2022-07-27 14:05:38 +00:00
frame = data[-1][list(self.sources)[0]]
if frame is not None:
# ADD FRAMETO VISION_CONSUMER QUEUE
self.vision_consumer.add_consumable({"frame": frame})
2022-06-22 15:18:29 +00:00
super()._get(emit=False)
2022-06-28 10:31:27 +00:00
def vision_consumer_work(self, consumable=None):
# VISION_CONSUMER TASK
2022-07-27 14:05:38 +00:00
if consumable is None:
2022-06-28 10:31:27 +00:00
return
detections = self.check_features(consumable["frame"])
results = self.process_detections(detections)
return {"detections": detections, "results": results}
def process_vision_consumed(self, data=None):
# print("VISION", str(int(QThread.currentThreadId())), flush=True)
if data is None:
return
# ADD VISION RETURNED FROM VISION_CONSUMER TO RENDER_CONSUMER QUEUE
data = data[-1][self.vision_consumer.name]
if data is not None:
# super()._get([{
# "frame": data["consumed"]["frame"],
# "detections": data["result"]["detections"],
# "results": data["result"]["results"],
# }])
self.render_consumer.add_consumable({
"frame": data["consumed"]["frame"],
"detections": data["result"]["detections"],
"results": data["result"]["results"],
})
def render_consumer_work(self, consumable=None):
# RENDER_CONSUMER TASK
if consumable is None:
return
render = QPixmap.fromImage(self.render(
consumable["frame"],
detections=consumable["detections"],
vision_results=consumable["results"],
mask=False,
offset=None,
))
return render
def process_render_consumed(self, data=None):
2022-06-22 15:18:29 +00:00
# print("VISION", str(int(QThread.currentThreadId())), flush=True)
if data is None:
2022-06-21 12:10:52 +00:00
return
2022-06-28 10:31:27 +00:00
# EMIT VISION AND RENDER RESULTS RETURNED FROM RENDER_CONSUMER
data = data[-1][self.render_consumer.name]
2022-06-22 15:18:29 +00:00
if data is not None:
super()._get([{
2022-06-28 10:31:27 +00:00
"frame": data["consumed"]["frame"],
"detections": data["consumed"]["detections"],
"results": data["consumed"]["results"],
"render": data["result"],
2022-06-22 15:18:29 +00:00
}])