forked from maxsitt/insect-detect
/
yolo_tracker_preview.py
176 lines (148 loc) · 7.66 KB
/
yolo_tracker_preview.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
#!/usr/bin/env python3
'''
Author: Maximilian Sittinger (https://github.com/maxsitt)
Website: https://maxsitt.github.io/insect-detect-docs/
License: GNU GPLv3 (https://choosealicense.com/licenses/gpl-3.0/)
This Python script does the following:
- run a custom YOLO object detection model (.blob format) on-device (Luxonis OAK)
-> inference on downscaled LQ frames (e.g. 320x320 px)
- use an object tracker to track detected objects and assign unique tracking IDs (on-device)
- show downscaled LQ frames + model/tracker output (bounding box, label, confidence,
tracking ID, tracking status) + fps in a new window (e.g. via X11 forwarding)
- optional argument:
"-log" print available Raspberry Pi memory, RPi CPU utilization + temperature,
OAK memory + CPU usage and OAK chip temperature to console
based on open source scripts available at https://github.com/luxonis
'''
import argparse
import json
import time
from pathlib import Path
import cv2
import depthai as dai
import numpy as np
# Define optional argument
parser = argparse.ArgumentParser()
parser.add_argument("-log", "--print_logs", action="store_true",
help="print RPi available memory, RPi CPU utilization + temperature, \
OAK memory + CPU usage and OAK chip temperature to console")
args = parser.parse_args()
if args.print_logs:
import psutil
from apscheduler.schedulers.background import BackgroundScheduler
from gpiozero import CPUTemperature
# Set file paths to the detection model and config JSON
MODEL_PATH = Path("insect-detect/models/yolov5n_320_openvino_2022.1_4shave.blob")
CONFIG_PATH = Path("insect-detect/models/json/yolov5_v7_320.json")
# Get detection model metadata from config JSON
with CONFIG_PATH.open(encoding="utf-8") as f:
config = json.load(f)
nn_config = config.get("nn_config", {})
nn_metadata = nn_config.get("NN_specific_metadata", {})
classes = nn_metadata.get("classes", {})
coordinates = nn_metadata.get("coordinates", {})
anchors = nn_metadata.get("anchors", {})
anchor_masks = nn_metadata.get("anchor_masks", {})
iou_threshold = nn_metadata.get("iou_threshold", {})
confidence_threshold = nn_metadata.get("confidence_threshold", {})
nn_mappings = config.get("mappings", {})
labels = nn_mappings.get("labels", {})
# Create depthai pipeline
pipeline = dai.Pipeline()
# Create and configure camera node
cam_rgb = pipeline.create(dai.node.ColorCamera)
#cam_rgb.setImageOrientation(dai.CameraImageOrientation.ROTATE_180_DEG)
cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
cam_rgb.setPreviewSize(320, 320) # downscaled LQ frames for model input
cam_rgb.setPreviewKeepAspectRatio(False) # "squeeze" frames (16:9) to square (1:1)
cam_rgb.setInterleaved(False) # planar layout
cam_rgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)
cam_rgb.setFps(47) # frames per second available for focus/exposure/model input
# Create detection network node and define input
nn = pipeline.create(dai.node.YoloDetectionNetwork)
cam_rgb.preview.link(nn.input) # downscaled LQ frames as model input
nn.input.setBlocking(False)
# Set detection model specific settings
nn.setBlobPath(MODEL_PATH)
nn.setNumClasses(classes)
nn.setCoordinateSize(coordinates)
nn.setAnchors(anchors)
nn.setAnchorMasks(anchor_masks)
nn.setIouThreshold(iou_threshold)
nn.setConfidenceThreshold(confidence_threshold)
nn.setNumInferenceThreads(2)
# Create and configure object tracker node and define inputs + outputs
tracker = pipeline.create(dai.node.ObjectTracker)
tracker.setTrackerType(dai.TrackerType.ZERO_TERM_IMAGELESS)
#tracker.setTrackerType(dai.TrackerType.SHORT_TERM_IMAGELESS) # better for low fps
tracker.setTrackerIdAssignmentPolicy(dai.TrackerIdAssignmentPolicy.UNIQUE_ID)
nn.passthrough.link(tracker.inputTrackerFrame)
nn.passthrough.link(tracker.inputDetectionFrame)
nn.out.link(tracker.inputDetections)
xout_rgb = pipeline.create(dai.node.XLinkOut)
xout_rgb.setStreamName("frame")
tracker.passthroughTrackerFrame.link(xout_rgb.input)
xout_tracker = pipeline.create(dai.node.XLinkOut)
xout_tracker.setStreamName("track")
tracker.out.link(xout_tracker.input)
# Define functions
def frame_norm(frame, bbox):
"""Convert relative bounding box coordinates (0-1) to pixel coordinates."""
norm_vals = np.full(len(bbox), frame.shape[0])
norm_vals[::2] = frame.shape[1]
return (np.clip(np.array(bbox), 0, 1) * norm_vals).astype(int)
def print_logs():
"""Print Raspberry Pi info to console."""
print(f"\nAvailable RPi memory: {round(psutil.virtual_memory().available / 1048576)} MB")
print(f"RPi CPU utilization: {round(psutil.cpu_percent(interval=None))} %")
print(f"RPi CPU temperature: {round(CPUTemperature().temperature)} °C\n")
# Connect to OAK device and start pipeline in USB2 mode
with dai.Device(pipeline, maxUsbSpeed=dai.UsbSpeed.HIGH) as device:
# Print RPi + OAK info to console every second
if args.print_logs:
scheduler = BackgroundScheduler()
scheduler.add_job(print_logs, "interval", seconds=1, id="log")
scheduler.start()
device.setLogLevel(dai.LogLevel.INFO)
device.setLogOutputLevel(dai.LogLevel.INFO)
# Create output queues to get the frames and tracklets + detections from the outputs defined above
q_frame = device.getOutputQueue(name="frame", maxSize=4, blocking=False)
q_track = device.getOutputQueue(name="track", maxSize=4, blocking=False)
# Create start_time and counter variable to measure fps
start_time = time.monotonic()
counter = 0
# Get LQ frames + tracker output (passthrough detections) and show in new window
while True:
if q_frame.has():
frame = q_frame.get().getCvFrame()
if q_track.has():
tracks = q_track.get().tracklets
counter += 1
fps = round(counter / (time.monotonic() - start_time), 2)
for track in tracks:
roi = track.roi.denormalize(frame.shape[1], frame.shape[0])
x1 = int(roi.topLeft().x)
y1 = int(roi.topLeft().y)
x2 = int(roi.bottomRight().x)
y2 = int(roi.bottomRight().y)
bbox = frame_norm(frame, (track.srcImgDetection.xmin, track.srcImgDetection.ymin,
track.srcImgDetection.xmax, track.srcImgDetection.ymax))
cv2.putText(frame, labels[track.srcImgDetection.label], (bbox[0], bbox[3] + 13),
cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)
cv2.putText(frame, f"{round(track.srcImgDetection.confidence, 2)}", (bbox[0], bbox[3] + 25),
cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)
cv2.putText(frame, f"ID:{track.id}", (bbox[0], bbox[3] + 40),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
cv2.putText(frame, track.status.name, (bbox[0], bbox[3] + 50),
cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255, 255, 255), 1)
cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 0, 255), 2) # model bbox
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 130), 1) # tracker bbox
cv2.putText(frame, f"fps: {fps}", (4, frame.shape[0] - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
cv2.imshow("tracker_preview", frame)
#print(f"fps: {fps}")
# streaming the frames via SSH (X11 forwarding) will slow down fps
# comment out "cv2.imshow()" and print fps to console for true fps
# Stop script and close window by pressing "Q"
if cv2.waitKey(1) == ord("q"):
break