Eduardo_Scripniciuc/st-ten-1
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

vision wip

Browse Source
This commit is contained in:
matteo porta 2022-06-27 11:03:26 +02:00
parent 4537c9fbc1
commit f25e0c2db6
3 changed files with 46 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
simulate.sh
View File

@ -17,7 +17,6 @@ python -B -u "./src/main.py" \
--auto-select \
--no-autotest \
--no-edgetpu \
--no-tflite \
--sim-camera \
--sim-modbus \
--sim-os-label-printer \
@ -28,6 +27,7 @@ $* 2> >(sed $'s/.*/\e[31m&\e[m/' >&2) # &
# --auto-login-user \
# --autotests-archive \
# --camera-edits \
# --no-tflite \
# --sim-archiver \
# --sim-serial-label-printer \
# --users-management \

57
src/components/vision.py
View File

@ -4,6 +4,7 @@ import traceback
from configparser import ConfigParser
from pathlib import Path
import cv2
import numpy
import numpy as np
import tensorflow as tf
@ -246,14 +247,14 @@ class Vision(Component):
if tf_mode is None and "edgetpu" in self.allowed_modes:
try:
# create tflite edgetpu interpreter
interpreter = make_interpreter(self.models_dir / model_name / f"{model_name}_edgetpu.tflite")
interpreter = make_interpreter(str(self.models_dir / model_name / f"{model_name}_edgetpu.tflite"))
tf_mode = "edgetpu"
except Exception:
self.log.exception(traceback.format_exc())
if tf_mode is None and "tflite_cpu" in self.allowed_modes:
try:
# create tflite cpu interpreter
interpreter = Interpreter(self.models_dir / model_name / f"{model_name}.tflite")
interpreter = Interpreter(str(self.models_dir / model_name / f"{model_name}.tflite"))
tf_mode = "tflite_cpu"
except Exception:
self.log.exception(traceback.format_exc())
@ -262,7 +263,7 @@ class Vision(Component):
if tf_mode is None and "normal" in self.allowed_modes:
try:
# create tensorflow model
model = tf.saved_model.load(self.models_dir / model_name)
model = tf.saved_model.load(str(self.models_dir / model_name))
tf_mode = "normal"
except Exception:
self.log.exception(traceback.format_exc())
@ -273,19 +274,18 @@ class Vision(Component):
self.model_name = model_name
if interpreter is not None:
# if there is a new tflite interpreter initialize it and the related values
interpreter.allocate_tensors()
interpreter.invoke() # warmup
self.tflite_input_details = interpreter.get_input_details()
self.tflite_output_details = interpreter.get_output_details()
self.inf_index = self.tflite_input_details[0]["index"]
self.inf_width = self.tflite_input_details[0]["shape"][2]
self.inf_height = self.tflite_input_details[0]["shape"][1]
self.inf_shape = self.tflite_input_details[0]["shape"]
# interpreter.resize_tensor_input(self.inf_index, self.inf_shape)
interpreter.allocate_tensors()
interpreter.invoke() # warmup
else:
self.tflite_input_details = None
self.tflite_output_details = None
self.inf_index = None
self.inf_width = None
self.inf_height = None
self.inf_shape = None
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:
@ -295,6 +295,8 @@ class Vision(Component):
self.loading_model_signal.emit({"status": "done"})
def check_features(self, image, lock=True):
if image.shape != self.inf_shape[1:3]:
image = cv2.resize(image, self.inf_shape[1:3], interpolation=cv2.INTER_LINEAR)
tensor = np.expand_dims(np.asarray(image), axis=0)
# Run inference
if lock:
@ -305,40 +307,41 @@ class Vision(Component):
objs = detect.get_objects(self.interpreter, self.detection_threshold, (1, 1))
if lock:
self.lock.unlock()
boxes = [(obj.bbox.ymin / self.inf_height, obj.bbox.xmin / self.inf_width, obj.bbox.ymax / self.inf_height, obj.bbox.xmax / self.inf_width) for obj in objs]
boxes = [(obj.bbox.ymin / self.inf_shape[1], obj.bbox.xmin / self.inf_shape[2], obj.bbox.ymax / self.inf_shape[1], obj.bbox.xmax / self.inf_shape[2]) for obj in objs]
classes = [obj.id + 1 for obj in objs]
scores = [obj.score for obj in objs]
detections = {
"detection_boxes": boxes,
"detection_classes": classes,
"detection_scores": scores,
"num_detections": 10,
"detection_boxes": [boxes],
"detection_classes": [classes],
"detection_scores": [scores],
}
else:
detections = self.model(tensor)
if lock:
self.lock.unlock()
# WARNING: results other than the ones related to tensor[-1] will be discarded
parsed_detections = []
for d_box, d_class, d_score, d_mask in zip(
detections["detection_boxes"][0],
detections["detection_classes"][0],
detections["detection_scores"][0],
detections["detection_masks"][0],
for d_box, d_class, d_score in zip( # , d_mask in zip(
detections["detection_boxes"][-1],
detections["detection_classes"][-1],
detections["detection_scores"][-1],
# detections["detection_masks"][-1],
):
if d_score < self.detection_threshold:
continue
center = self.get_center(d_box.numpy().tolist())
box = list(d_box)
# box = d_box.numpy().tolist()
center = self.get_center(box)
detection = {
"class": self.category_index[int(d_class)]["name"],
"color": self.category_index[int(d_class)]["color_qt"],
"score": d_score.numpy().tolist(),
"mask": d_mask.numpy().tolist(),
"box": d_box.numpy().tolist(),
"class": self.category_index[int(d_class)],
"score": float(d_score),
# "score": d_score.numpy().tolist(),
# "mask": d_mask.numpy().tolist(),
"box": box,
"center": center,
"pos_rel_mm": self.get_pos_rel_mm(center),
}
parsed_detections.append(detection)
return {"result": parsed_detections, "ok": False}
return {"result": parsed_detections, "shape": tensor[-1].shape, "ok": False, "tensor": tensor[-1]}
def _get(self, data):
# print("VISION", str(int(QThread.currentThreadId())), flush=True)

16
src/ui/test_vision/test_vision.py
View File

@ -100,10 +100,15 @@ class Test_Vision(Widget):
data = data["vision"]
frame = data.get("frame", None)
vision = data.get("vision", None)
if vision is not None:
tensor = vision.get("tensor", None)
else:
tensor = None
rendered = data.get("rendered", None)
self.last_vision = {
"time": time,
"frame": frame,
"tensor": tensor,
"vision": vision,
"rendered": rendered,
}
@ -122,12 +127,13 @@ class Test_Vision(Widget):
self.visualize_vision(
time=time,
frame=frame,
tensor=tensor,
vision=vision,
rendered=rendered,
overridden=override,
)
def visualize_vision(self, time=None, frame=None, vision=None, rendered=None, overridden=False):
def visualize_vision(self, time=None, frame=None, tensor=None, vision=None, rendered=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"])
@ -155,6 +161,14 @@ class Test_Vision(Widget):
frame.shape[2] * frame.shape[1], # width * channels
QImage.Format_RGB888
))
elif tensor is not None:
self.img = QPixmap.fromImage(QImage(
tensor.data,
tensor.shape[1], # width
tensor.shape[0], # height
tensor.shape[2] * tensor.shape[1], # width * channels
QImage.Format_RGB888
))
else:
self.img = self.status_imgs_full[None]
self.resizeEvent()