diff --git a/py/src/config.json b/py/src/config.json new file mode 100644 index 00000000..420250c6 --- /dev/null +++ b/py/src/config.json @@ -0,0 +1,7 @@ +{ + "device_id": "ai-spy", + "server_ip": "10.44.67.2", + "stream_port": 8000, + "model_path": "./model/edgetpu.tflite", + "label_path": "./model/labelmap.pbtxt" +} diff --git a/py/src/main.py b/py/src/main.py index ee164c4c..3dbe205b 100644 --- a/py/src/main.py +++ b/py/src/main.py @@ -1,7 +1,7 @@ from config.Config import ConfigStore, LocalConfig, RemoteConfig from config.ConfigSource import FileConfigSource, NTConfigSource from pipeline.Capture import CVCapture -from pipeline.Detector import Detector, draw_bbox +from pipeline.Detector import Detector from output.StreamServer import StreamServer import time @@ -40,5 +40,5 @@ objs = detector.run_inference(image) - image = draw_bbox(image, objs) + image = detector.draw_bbox(image, objs) stream_server.set_frame(image) diff --git a/py/src/pipeline/Detector.py b/py/src/pipeline/Detector.py index d85f9512..26c7f2a1 100644 --- a/py/src/pipeline/Detector.py +++ b/py/src/pipeline/Detector.py @@ -11,15 +11,6 @@ import cv2 as cv -def draw_bbox(image: cv.Mat, objs: list[Object]) -> cv.Mat: - for obj in objs: - cv.rectangle(image, - (obj.bbox.xmin.item(), obj.bbox.ymin.item()), - (obj.bbox.xmax.item(), obj.bbox.ymax.item()), - (0, 255, 0)) - return image - - class Detector: """Runs a tflite model on the coral to detect objects""" _labels: Dict[int, str] @@ -46,3 +37,20 @@ def run_inference(self, image: cv.Mat) -> list[Object]: run_inference(self._interpreter, image) objs = get_objects(self._interpreter, self._config.remote_config.detection_threshold)[:self._config.remote_config.max_targets] return objs + + def draw_bbox(self, image: cv.Mat, objs: list[Object]) -> cv.Mat: + height, width, channels = image.shape + scale_x, scale_y = width / self._inference_size[0], height / self._inference_size[1] + for obj in objs: + bbox = obj.bbox.scale(scale_x, scale_y) + x0, y0 = int(bbox.xmin), int(bbox.ymin) + x1, y1 = int(bbox.xmax), int(bbox.ymax) + + percent = int(100 * obj.score) + label = '{}% {}'.format(percent, self._labels.get(obj.id, obj.id)) + + image = cv.rectangle(image, (x0, y0), (x1, y1), (0, 255, 0), 2) + image = cv.putText(image, label, (x0, y0+30), + cv.FONT_HERSHEY_SIMPLEX, 1.0, (255, 0, 0), 2) + + return image