From 2e9d208ad860f26b243d75a15369c2a7c665d8d1 Mon Sep 17 00:00:00 2001 From: Tommy Date: Mon, 26 Jun 2023 01:12:30 +0200 Subject: [PATCH] Add mask detection --- robothub-oak/Mask Detection/Dockerfile | 6 ++ robothub-oak/Mask Detection/app.py | 113 ++++++++++++++++++++++ robothub-oak/Mask Detection/robotapp.toml | 73 ++++++++++++++ 3 files changed, 192 insertions(+) create mode 100644 robothub-oak/Mask Detection/Dockerfile create mode 100644 robothub-oak/Mask Detection/app.py create mode 100644 robothub-oak/Mask Detection/robotapp.toml diff --git a/robothub-oak/Mask Detection/Dockerfile b/robothub-oak/Mask Detection/Dockerfile new file mode 100644 index 0000000..d9a748f --- /dev/null +++ b/robothub-oak/Mask Detection/Dockerfile @@ -0,0 +1,6 @@ +FROM ghcr.io/luxonis/robothub-app-v2:2023.108.0914-regular + +RUN apt-get update && apt-get install -y git +RUN pip3 uninstall -y depthai_sdk robothub_oak +RUN pip3 install -U git+https://github.com/luxonis/depthai.git@develop#subdirectory=depthai_sdk +RUN pip3 install -U git+https://github.com/luxonis/robothub-oak.git@develop \ No newline at end of file diff --git a/robothub-oak/Mask Detection/app.py b/robothub-oak/Mask Detection/app.py new file mode 100644 index 0000000..626e133 --- /dev/null +++ b/robothub-oak/Mask Detection/app.py @@ -0,0 +1,113 @@ +from robothub_oak.manager import DEVICE_MANAGER +from collections import defaultdict +from depthai_sdk import Visualizer +from robothub_oak.packets import DetectionPacket +import numpy as np +import robothub +import warnings +import time +import json +import cv2 + +def log_softmax(x): + e_x = np.exp(x - np.max(x)) + return np.log(e_x / e_x.sum()) + +class ExampleApplication(robothub.RobotHubApplication): + def __init__(self): + super().__init__() + + self.streams = defaultdict(lambda: defaultdict(robothub.StreamHandle)) + self.detections = [] + self.old_detections = [] + + def on_start(self): + mask_detections = robothub.CONFIGURATION['mask_detections'] + mask_detections_mode = robothub.CONFIGURATION['mask_detections_mode'] + helmet_detections = robothub.CONFIGURATION['helmet_detections'] + helmet_detections_mode = robothub.CONFIGURATION['helmet_detections_mode'] + vest_detections = robothub.CONFIGURATION['vest_detections'] + vest_detections_mode = robothub.CONFIGURATION['vest_detections_mode'] + + print(f""" +Starting with: +Mask detections {"enabled" if mask_detections else "disabled"}, mode {mask_detections_mode} +Helmet detections {"enabled" if helmet_detections else "disabled"}, mode {helmet_detections_mode} +Vest detections {"enabled" if vest_detections else "disabled"}, mode {vest_detections_mode} +""") + devices = DEVICE_MANAGER.get_all_devices() + for device in devices: + color_resolution = '1080p' + + color = device.get_camera('color', resolution=color_resolution, fps=30) + color.stream_to_hub(name=f'Color stream {device.get_device_name()}', unique_key=f'Color {device.id}') + + if mask_detections: + face_nn = device.create_neural_network('face-detection-retail-0004', color) + + face_nn.configure(resize_mode='crop') + + mask_nn = device.create_neural_network('sbd_mask_classification_224x224', face_nn) + + mask_nn.add_callback(self.on_mask, output_type='passthrough') + + face_nn.add_callback(self.on_face, output_type='encoded') + + stream_handle = robothub.STREAMS.create_video(device.mxid, f'NN {device.id}', f'Detection stream {device.get_device_name()}') + self.streams[device.mxid]['NN'] = stream_handle + + def on_mask(self, packet): + self.old_detections = self.detections + self.detections = [] + + for det, rec in zip(packet.detections, packet.nn_data): + index = np.argmax(log_softmax(rec.getFirstLayerFp16())) + self.detections.append(index) + + + def on_face(self, packet: DetectionPacket): + vis: Visualizer = packet.visualizer + + try: + detections = vis.objects[0] + for i, mask in enumerate(self.detections): + detections.labels[i] = "Mask" if mask else "No mask" + detections.colors[i] = (0, 255, 0) if mask else (255, 0, 0) + except: + pass + + metadata = json.loads(packet.visualizer.serialize()) + timestamp = int(time.time() * 1_000) + frame_bytes = bytes(packet.msg.getData()) + self.streams[packet.device.mxid]['NN'].publish_video_data(frame_bytes, timestamp, metadata) + + def upload_detection(packet, metadata, tags): + try: + frame_bytes = cv2.imencode('.jpg', packet.depthai_sdk_packet.frame)[1].tobytes() + + event = robothub.EVENTS.prepare() + event.add_frame(frame_bytes, packet.device.mxid, metadata=metadata) + event.add_tags(tags) + robothub.EVENTS.upload(event) + + except Exception as e: + warnings.warn(f'Could not upload detection with error: {e}') + + mode = robothub.CONFIGURATION['mask_detections_mode'] + for old, new in zip(self.old_detections, self.detections): + if mode == 'not_wearing': + if old == 1 and new == 0: + upload_detection(packet, metadata, ["Mask", "Not wearing"]) + elif mode == 'wearing': + if old == 0 and new == 1: + upload_detection(packet, metadata, ["Mask", "Wearing"]) + elif mode == 'change': + if old != new: + upload_detection(packet, metadata, ["Mask", "Change"]) + + + def start_execution(self): + DEVICE_MANAGER.start() + + def on_stop(self): + DEVICE_MANAGER.stop() diff --git a/robothub-oak/Mask Detection/robotapp.toml b/robothub-oak/Mask Detection/robotapp.toml new file mode 100644 index 0000000..f02d933 --- /dev/null +++ b/robothub-oak/Mask Detection/robotapp.toml @@ -0,0 +1,73 @@ +config_version = "2.0" + +[info] +name = "Mask detection" +description = "Detects if people are wearing masks" + +[runtime] +application = "app.py#ExampleApplication" +workdir = "/app" + +[runtime.runs_on] +type = "custom" +dockerfile = "Dockerfile" + +[[configuration]] +field = "boolean" +key = "mask_detections" +label = "Enable mask detections" +initial_value = true + +[[configuration]] +field = "choice" +key = "mask_detections_mode" +label = "Send mask detections when" +style = "dropdown" +allow_multiple = false +options = [ + { key = "not_wearing", label = "Not wearing", default = true }, + { key = "wearing", label = "Wearing" }, + { key = "change", label = "Change" }, +] + +[[configuration]] +visual = "divider" + +[[configuration]] +field = "boolean" +key = "helmet_detections" +label = "Enable helmet detections" +initial_value = false + +[[configuration]] +field = "choice" +key = "helmet_detections_mode" +label = "Send helmet detections when" +style = "dropdown" +allow_multiple = false +options = [ + { key = "not_wearing", label = "Not wearing", default = true }, + { key = "wearing", label = "Wearing" }, + { key = "change", label = "Change" }, +] + +[[configuration]] +visual = "divider" + +[[configuration]] +field = "boolean" +key = "vest_detections" +label = "Enable vest detections" +initial_value = false + +[[configuration]] +field = "choice" +key = "vest_detections_mode" +label = "Send vest detections when" +style = "dropdown" +allow_multiple = false +options = [ + { key = "not_wearing", label = "Not wearing", default = true }, + { key = "wearing", label = "Wearing" }, + { key = "change", label = "Change" }, +] \ No newline at end of file