353 lines
11 KiB
Python
353 lines
11 KiB
Python
|
import cv2
|
||
|
|
import depthai as dai
|
||
|
|
import blobconverter
|
||
|
|
import numpy as np
|
||
|
|
import time
|
||
|
|
import math
|
||
|
|
|
||
|
|
|
||
|
|
LABELS = [
|
||
|
|
"background", "aeroplane", "bicycle", "bird", "boat", "bottle", "bus",
|
||
|
|
"car", "cat", "chair", "cow", "diningtable", "dog", "horse",
|
||
|
|
"motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
# =========================
|
||
|
|
# CONFIG SEGURANÇA
|
||
|
|
# =========================
|
||
|
|
|
||
|
|
MIN_DEPTH_MM = 300
|
||
|
|
MAX_DEPTH_MM = 3500
|
||
|
|
DANGER_DEPTH_MM = 1600
|
||
|
|
|
||
|
|
MIN_BLOB_AREA_PX = 120
|
||
|
|
TRACK_MAX_DIST = 55
|
||
|
|
MAX_LOST_FRAMES = 10
|
||
|
|
|
||
|
|
ROI_TOP = 0.25
|
||
|
|
ROI_BOTTOM = 0.95
|
||
|
|
ROI_LEFT = 0.12
|
||
|
|
ROI_RIGHT = 0.88
|
||
|
|
|
||
|
|
AI_CONFIDENCE = 0.5
|
||
|
|
AI_DANGER_CLASSES = {
|
||
|
|
"person", "bicycle", "motorbike", "car", "bus", "dog", "cat", "cow", "horse", "sheep"
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
class SimpleBlobTracker:
|
||
|
|
def __init__(self):
|
||
|
|
self.next_id = 1
|
||
|
|
self.tracks = {}
|
||
|
|
|
||
|
|
def update(self, detections):
|
||
|
|
updated = []
|
||
|
|
used_tracks = set()
|
||
|
|
|
||
|
|
for det in detections:
|
||
|
|
best_id = None
|
||
|
|
best_dist = 999999
|
||
|
|
|
||
|
|
for tid, tr in self.tracks.items():
|
||
|
|
if tid in used_tracks:
|
||
|
|
continue
|
||
|
|
|
||
|
|
dx = det["cx"] - tr["cx"]
|
||
|
|
dy = det["cy"] - tr["cy"]
|
||
|
|
dz = (det["z_mm"] - tr["z_mm"]) / 30.0
|
||
|
|
dist = math.sqrt(dx * dx + dy * dy + dz * dz)
|
||
|
|
|
||
|
|
if dist < best_dist:
|
||
|
|
best_dist = dist
|
||
|
|
best_id = tid
|
||
|
|
|
||
|
|
if best_id is not None and best_dist < TRACK_MAX_DIST:
|
||
|
|
tid = best_id
|
||
|
|
used_tracks.add(tid)
|
||
|
|
self.tracks[tid].update(det)
|
||
|
|
self.tracks[tid]["lost"] = 0
|
||
|
|
self.tracks[tid]["age"] += 1
|
||
|
|
else:
|
||
|
|
tid = self.next_id
|
||
|
|
self.next_id += 1
|
||
|
|
self.tracks[tid] = dict(det)
|
||
|
|
self.tracks[tid]["lost"] = 0
|
||
|
|
self.tracks[tid]["age"] = 1
|
||
|
|
|
||
|
|
out = dict(self.tracks[tid])
|
||
|
|
out["id"] = tid
|
||
|
|
updated.append(out)
|
||
|
|
|
||
|
|
for tid in list(self.tracks.keys()):
|
||
|
|
if tid not in used_tracks and all(d.get("id") != tid for d in updated):
|
||
|
|
self.tracks[tid]["lost"] += 1
|
||
|
|
if self.tracks[tid]["lost"] > MAX_LOST_FRAMES:
|
||
|
|
del self.tracks[tid]
|
||
|
|
|
||
|
|
return updated
|
||
|
|
|
||
|
|
|
||
|
|
# =========================
|
||
|
|
# PIPELINE OAK
|
||
|
|
# =========================
|
||
|
|
|
||
|
|
pipeline = dai.Pipeline()
|
||
|
|
|
||
|
|
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_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)
|
||
|
|
|
||
|
|
stereo = pipeline.create(dai.node.StereoDepth)
|
||
|
|
stereo.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.DEFAULT)
|
||
|
|
stereo.setDepthAlign(dai.CameraBoardSocket.CAM_A)
|
||
|
|
stereo.setSubpixel(True)
|
||
|
|
stereo.setLeftRightCheck(True)
|
||
|
|
|
||
|
|
mono_left.out.link(stereo.left)
|
||
|
|
mono_right.out.link(stereo.right)
|
||
|
|
|
||
|
|
# IA espacial
|
||
|
|
detection = pipeline.create(dai.node.MobileNetSpatialDetectionNetwork)
|
||
|
|
detection.setBlobPath(blobconverter.from_zoo(
|
||
|
|
name="mobilenet-ssd",
|
||
|
|
shaves=3,
|
||
|
|
version="2021.4"
|
||
|
|
))
|
||
|
|
detection.setConfidenceThreshold(AI_CONFIDENCE)
|
||
|
|
detection.input.setBlocking(False)
|
||
|
|
detection.setBoundingBoxScaleFactor(0.5)
|
||
|
|
detection.setDepthLowerThreshold(MIN_DEPTH_MM)
|
||
|
|
detection.setDepthUpperThreshold(8000)
|
||
|
|
|
||
|
|
cam_rgb.preview.link(detection.input)
|
||
|
|
stereo.depth.link(detection.inputDepth)
|
||
|
|
|
||
|
|
# Tracker oficial para IA
|
||
|
|
tracker_ai = pipeline.create(dai.node.ObjectTracker)
|
||
|
|
tracker_ai.setTrackerType(dai.TrackerType.ZERO_TERM_COLOR_HISTOGRAM)
|
||
|
|
tracker_ai.setTrackerIdAssignmentPolicy(dai.TrackerIdAssignmentPolicy.SMALLEST_ID)
|
||
|
|
|
||
|
|
detection.passthrough.link(tracker_ai.inputTrackerFrame)
|
||
|
|
detection.passthrough.link(tracker_ai.inputDetectionFrame)
|
||
|
|
detection.out.link(tracker_ai.inputDetections)
|
||
|
|
|
||
|
|
# Outputs
|
||
|
|
xout_rgb = pipeline.create(dai.node.XLinkOut)
|
||
|
|
xout_depth = pipeline.create(dai.node.XLinkOut)
|
||
|
|
xout_ai = pipeline.create(dai.node.XLinkOut)
|
||
|
|
|
||
|
|
xout_rgb.setStreamName("rgb")
|
||
|
|
xout_depth.setStreamName("depth")
|
||
|
|
xout_ai.setStreamName("ai_tracklets")
|
||
|
|
|
||
|
|
cam_rgb.preview.link(xout_rgb.input)
|
||
|
|
#stereo.depth.link(xout_depth.input)
|
||
|
|
tracker_ai.out.link(xout_ai.input)
|
||
|
|
|
||
|
|
|
||
|
|
# =========================
|
||
|
|
# LOOP
|
||
|
|
# =========================
|
||
|
|
|
||
|
|
blob_tracker = SimpleBlobTracker()
|
||
|
|
|
||
|
|
with dai.Device(pipeline) as device:
|
||
|
|
q_rgb = device.getOutputQueue("rgb", maxSize=1, blocking=True)
|
||
|
|
q_depth = device.getOutputQueue("depth", maxSize=1, blocking=False)
|
||
|
|
q_ai = device.getOutputQueue("ai_tracklets", maxSize=1, blocking=False)
|
||
|
|
|
||
|
|
last_depth = None
|
||
|
|
last_ai_tracklets = []
|
||
|
|
|
||
|
|
last = time.time()
|
||
|
|
fps = 0
|
||
|
|
|
||
|
|
while True:
|
||
|
|
in_rgb = q_rgb.get()
|
||
|
|
frame = in_rgb.getCvFrame()
|
||
|
|
|
||
|
|
in_depth = q_depth.tryGet()
|
||
|
|
if in_depth is not None:
|
||
|
|
last_depth = in_depth.getFrame()
|
||
|
|
|
||
|
|
in_ai = q_ai.tryGet()
|
||
|
|
if in_ai is not None:
|
||
|
|
last_ai_tracklets = in_ai.tracklets
|
||
|
|
|
||
|
|
depth_ok = last_depth is not None
|
||
|
|
|
||
|
|
if depth_ok:
|
||
|
|
depth = last_depth
|
||
|
|
|
||
|
|
if depth.shape[:2] != frame.shape[:2]:
|
||
|
|
depth = cv2.resize(depth, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_NEAREST)
|
||
|
|
else:
|
||
|
|
depth = None
|
||
|
|
|
||
|
|
depth = last_depth
|
||
|
|
ai_tracklets = last_ai_tracklets
|
||
|
|
|
||
|
|
h, w = frame.shape[:2]
|
||
|
|
|
||
|
|
now = time.time()
|
||
|
|
fps = 0.9 * fps + 0.1 * (1 / max(now - last, 1e-6))
|
||
|
|
last = now
|
||
|
|
|
||
|
|
danger_ai = False
|
||
|
|
danger_blob = False
|
||
|
|
|
||
|
|
# =========================
|
||
|
|
# 1) IA + DEPTH + TRACKER
|
||
|
|
# =========================
|
||
|
|
|
||
|
|
for t in ai_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)
|
||
|
|
z_mm = t.spatialCoordinates.z
|
||
|
|
z_m = z_mm / 1000.0
|
||
|
|
|
||
|
|
is_danger_class = label in AI_DANGER_CLASSES
|
||
|
|
is_close = MIN_DEPTH_MM < z_mm < DANGER_DEPTH_MM
|
||
|
|
|
||
|
|
if is_danger_class and is_close:
|
||
|
|
danger_ai = True
|
||
|
|
|
||
|
|
color = (0, 0, 255) if is_danger_class and is_close else (0, 180, 255)
|
||
|
|
|
||
|
|
txt = f"AI ID {t.id} | {label} | Z={z_m:.2f}m | {t.status.name}"
|
||
|
|
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
|
||
|
|
cv2.putText(frame, txt, (x1, max(20, y1 - 8)),
|
||
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.42, color, 1)
|
||
|
|
|
||
|
|
# =========================
|
||
|
|
# 2) DEPTH BLOB SEM IA
|
||
|
|
# =========================
|
||
|
|
danger_blob = False
|
||
|
|
tracked_blobs = []
|
||
|
|
mask = None
|
||
|
|
|
||
|
|
if depth_ok:
|
||
|
|
x1_roi = int(w * ROI_LEFT)
|
||
|
|
x2_roi = int(w * ROI_RIGHT)
|
||
|
|
y1_roi = int(h * ROI_TOP)
|
||
|
|
y2_roi = int(h * ROI_BOTTOM)
|
||
|
|
|
||
|
|
roi_depth = depth[y1_roi:y2_roi, x1_roi:x2_roi]
|
||
|
|
|
||
|
|
mask = np.zeros_like(roi_depth, dtype=np.uint8)
|
||
|
|
mask[(roi_depth > MIN_DEPTH_MM) & (roi_depth < MAX_DEPTH_MM)] = 255
|
||
|
|
|
||
|
|
kernel = np.ones((5, 5), np.uint8)
|
||
|
|
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
|
||
|
|
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
|
||
|
|
|
||
|
|
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
|
|
||
|
|
blob_detections = []
|
||
|
|
|
||
|
|
for cnt in contours:
|
||
|
|
area = cv2.contourArea(cnt)
|
||
|
|
if area < MIN_BLOB_AREA_PX:
|
||
|
|
continue
|
||
|
|
|
||
|
|
x, y, bw, bh = cv2.boundingRect(cnt)
|
||
|
|
gx = x + x1_roi
|
||
|
|
gy = y + y1_roi
|
||
|
|
gcx = gx + bw // 2
|
||
|
|
gcy = gy + bh // 2
|
||
|
|
|
||
|
|
blob_depth = roi_depth[y:y + bh, x:x + bw]
|
||
|
|
valid = blob_depth[(blob_depth > MIN_DEPTH_MM) & (blob_depth < MAX_DEPTH_MM)]
|
||
|
|
|
||
|
|
if len(valid) < 50:
|
||
|
|
continue
|
||
|
|
|
||
|
|
z_mm = float(np.median(valid))
|
||
|
|
|
||
|
|
blob_detections.append({
|
||
|
|
"bbox": (gx, gy, gx + bw, gy + bh),
|
||
|
|
"cx": gcx,
|
||
|
|
"cy": gcy,
|
||
|
|
"z_mm": z_mm,
|
||
|
|
"area": area
|
||
|
|
})
|
||
|
|
|
||
|
|
tracked_blobs = blob_tracker.update(blob_detections)
|
||
|
|
|
||
|
|
cv2.rectangle(frame, (x1_roi, y1_roi), (x2_roi, y2_roi), (255, 255, 0), 1)
|
||
|
|
|
||
|
|
for obj in tracked_blobs:
|
||
|
|
x1, y1, x2, y2 = obj["bbox"]
|
||
|
|
z_mm = obj["z_mm"]
|
||
|
|
z_m = z_mm / 1000.0
|
||
|
|
center_percent = obj["cx"] / w * 100.0
|
||
|
|
|
||
|
|
is_close = z_mm < DANGER_DEPTH_MM
|
||
|
|
if is_close:
|
||
|
|
danger_blob = True
|
||
|
|
|
||
|
|
color = (0, 0, 255) if is_close else (0, 255, 0)
|
||
|
|
|
||
|
|
txt = f"DEPTH ID {obj['id']} | Z={z_m:.2f}m | X={center_percent:.0f}% | area={int(obj['area'])}"
|
||
|
|
cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
|
||
|
|
cv2.circle(frame, (obj["cx"], obj["cy"]), 4, (255, 0, 255), -1)
|
||
|
|
cv2.putText(frame, txt, (x1, min(h - 10, y2 + 16)), cv2.FONT_HERSHEY_SIMPLEX, 0.42, color, 1)
|
||
|
|
else:
|
||
|
|
cv2.putText(frame, "Depth: aguardando/indisponivel", (10, 115), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 255, 255), 1)
|
||
|
|
|
||
|
|
# =========================
|
||
|
|
# 3) DECISÃO FINAL
|
||
|
|
# =========================
|
||
|
|
|
||
|
|
danger_final = danger_ai or danger_blob
|
||
|
|
|
||
|
|
cv2.putText(frame, f"FPS: {fps:.1f}", (10, 25),
|
||
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
|
||
|
|
|
||
|
|
cv2.putText(frame, f"AI danger: {danger_ai} | Depth danger: {danger_blob}",
|
||
|
|
(10, 52), cv2.FONT_HERSHEY_SIMPLEX, 0.55,
|
||
|
|
(255, 255, 255), 1)
|
||
|
|
|
||
|
|
if danger_final:
|
||
|
|
cv2.putText(frame, "PARAR: risco detectado", (10, 85),
|
||
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.85, (0, 0, 255), 2)
|
||
|
|
else:
|
||
|
|
cv2.putText(frame, "Livre", (10, 85),
|
||
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.85, (0, 255, 0), 2)
|
||
|
|
|
||
|
|
if depth_ok:
|
||
|
|
depth_vis = depth.copy()
|
||
|
|
depth_vis[depth_vis == 0] = MAX_DEPTH_MM
|
||
|
|
depth_vis = np.clip(depth_vis, MIN_DEPTH_MM, MAX_DEPTH_MM)
|
||
|
|
depth_vis = ((MAX_DEPTH_MM - depth_vis) / (MAX_DEPTH_MM - MIN_DEPTH_MM) * 255).astype(np.uint8)
|
||
|
|
depth_vis = cv2.applyColorMap(depth_vis, cv2.COLORMAP_JET)
|
||
|
|
cv2.imshow("Depth", depth_vis)
|
||
|
|
|
||
|
|
if mask is not None:
|
||
|
|
cv2.imshow("Depth Blob Mask", mask)
|
||
|
|
|
||
|
|
cv2.imshow("OAK-D Lite - Hybrid Safety Tracker", frame)
|
||
|
|
|
||
|
|
|
||
|
|
key = cv2.waitKey(1)
|
||
|
|
if key == ord("q") or key == 27:
|
||
|
|
break
|
||
|
|
|
||
|
|
cv2.destroyAllWindows()
|