137 lines
4.2 KiB
Python
137 lines
4.2 KiB
Python
|
import cv2
|
||
|
|
import depthai as dai
|
||
|
|
import blobconverter
|
||
|
|
import time
|
||
|
|
|
||
|
|
# COCO labels do MobileNet-SSD
|
||
|
|
LABELS = [
|
||
|
|
"background", "aeroplane", "bicycle", "bird", "boat", "bottle", "bus",
|
||
|
|
"car", "cat", "chair", "cow", "diningtable", "dog", "horse",
|
||
|
|
"motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"
|
||
|
|
]
|
||
|
|
|
||
|
|
pipeline = dai.Pipeline()
|
||
|
|
|
||
|
|
# RGB
|
||
|
|
cam_rgb = pipeline.create(dai.node.ColorCamera)
|
||
|
|
cam_rgb.setPreviewSize(300, 300)
|
||
|
|
cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
|
||
|
|
cam_rgb.setInterleaved(False)
|
||
|
|
cam_rgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)
|
||
|
|
cam_rgb.setFps(30)
|
||
|
|
|
||
|
|
# Mono stereo
|
||
|
|
mono_left = pipeline.create(dai.node.MonoCamera)
|
||
|
|
mono_right = pipeline.create(dai.node.MonoCamera)
|
||
|
|
|
||
|
|
mono_left.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)
|
||
|
|
mono_right.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)
|
||
|
|
|
||
|
|
mono_left.setBoardSocket(dai.CameraBoardSocket.CAM_B)
|
||
|
|
mono_right.setBoardSocket(dai.CameraBoardSocket.CAM_C)
|
||
|
|
|
||
|
|
# Depth
|
||
|
|
stereo = pipeline.create(dai.node.StereoDepth)
|
||
|
|
stereo.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.DEFAULT)
|
||
|
|
stereo.setDepthAlign(dai.CameraBoardSocket.CAM_A)
|
||
|
|
stereo.setSubpixel(True)
|
||
|
|
|
||
|
|
mono_left.out.link(stereo.left)
|
||
|
|
mono_right.out.link(stereo.right)
|
||
|
|
|
||
|
|
# Spatial Detection Network - MobileNet SSD
|
||
|
|
detection = pipeline.create(dai.node.MobileNetSpatialDetectionNetwork)
|
||
|
|
|
||
|
|
detection.setBlobPath(blobconverter.from_zoo(
|
||
|
|
name="mobilenet-ssd",
|
||
|
|
shaves=3,
|
||
|
|
version="2021.4"
|
||
|
|
))
|
||
|
|
|
||
|
|
detection.setConfidenceThreshold(0.5)
|
||
|
|
detection.input.setBlocking(False)
|
||
|
|
detection.setBoundingBoxScaleFactor(0.5)
|
||
|
|
detection.setDepthLowerThreshold(200)
|
||
|
|
detection.setDepthUpperThreshold(8000)
|
||
|
|
|
||
|
|
cam_rgb.preview.link(detection.input)
|
||
|
|
stereo.depth.link(detection.inputDepth)
|
||
|
|
|
||
|
|
# Tracker
|
||
|
|
tracker = pipeline.create(dai.node.ObjectTracker)
|
||
|
|
|
||
|
|
# Para obstáculo geral, rastreia tudo detectado.
|
||
|
|
# Para só pessoas: tracker.setDetectionLabelsToTrack([15])
|
||
|
|
tracker.setTrackerType(dai.TrackerType.ZERO_TERM_COLOR_HISTOGRAM)
|
||
|
|
tracker.setTrackerIdAssignmentPolicy(dai.TrackerIdAssignmentPolicy.SMALLEST_ID)
|
||
|
|
|
||
|
|
detection.passthrough.link(tracker.inputTrackerFrame)
|
||
|
|
detection.passthrough.link(tracker.inputDetectionFrame)
|
||
|
|
detection.out.link(tracker.inputDetections)
|
||
|
|
|
||
|
|
# Outputs
|
||
|
|
xout_rgb = pipeline.create(dai.node.XLinkOut)
|
||
|
|
xout_track = pipeline.create(dai.node.XLinkOut)
|
||
|
|
|
||
|
|
xout_rgb.setStreamName("rgb")
|
||
|
|
xout_track.setStreamName("tracklets")
|
||
|
|
|
||
|
|
tracker.passthroughTrackerFrame.link(xout_rgb.input)
|
||
|
|
tracker.out.link(xout_track.input)
|
||
|
|
|
||
|
|
with dai.Device(pipeline) as device:
|
||
|
|
q_rgb = device.getOutputQueue("rgb", maxSize=4, blocking=False)
|
||
|
|
q_track = device.getOutputQueue("tracklets", maxSize=4, blocking=False)
|
||
|
|
|
||
|
|
last = time.time()
|
||
|
|
fps = 0
|
||
|
|
|
||
|
|
while True:
|
||
|
|
frame = q_rgb.get().getCvFrame()
|
||
|
|
tracklets = q_track.get().tracklets
|
||
|
|
|
||
|
|
now = time.time()
|
||
|
|
fps = 0.9 * fps + 0.1 * (1 / max(now - last, 1e-6))
|
||
|
|
last = now
|
||
|
|
|
||
|
|
h, w = frame.shape[:2]
|
||
|
|
|
||
|
|
perigo = False
|
||
|
|
|
||
|
|
for t in tracklets:
|
||
|
|
roi = t.roi.denormalize(w, h)
|
||
|
|
x1 = int(roi.topLeft().x)
|
||
|
|
y1 = int(roi.topLeft().y)
|
||
|
|
x2 = int(roi.bottomRight().x)
|
||
|
|
y2 = int(roi.bottomRight().y)
|
||
|
|
|
||
|
|
label = LABELS[t.label] if t.label < len(LABELS) else str(t.label)
|
||
|
|
|
||
|
|
x_mm = t.spatialCoordinates.x
|
||
|
|
y_mm = t.spatialCoordinates.y
|
||
|
|
z_mm = t.spatialCoordinates.z
|
||
|
|
|
||
|
|
dist_m = z_mm / 1000.0
|
||
|
|
|
||
|
|
if dist_m < 1.5:
|
||
|
|
perigo = True
|
||
|
|
|
||
|
|
texto = f"ID {t.id} | {label} | {t.status.name} | Z={dist_m:.2f}m"
|
||
|
|
|
||
|
|
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
||
|
|
cv2.putText(frame, texto, (x1, max(20, y1 - 8)), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255), 1)
|
||
|
|
|
||
|
|
cv2.circle(frame, ((x1 + x2) // 2, (y1 + y2) // 2), 4, (0, 255, 255), -1)
|
||
|
|
|
||
|
|
cv2.putText(frame, f"FPS: {fps:.1f}", (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
|
||
|
|
|
||
|
|
if perigo:
|
||
|
|
cv2.putText(frame, "PERIGO: objeto perto - PARAR", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
|
||
|
|
|
||
|
|
cv2.imshow("OAK-D Lite Spatial Object Tracker", frame)
|
||
|
|
|
||
|
|
key = cv2.waitKey(1)
|
||
|
|
if key == ord("q") or key == 27:
|
||
|
|
break
|
||
|
|
|
||
|
|
cv2.destroyAllWindows()
|