#!/usr/bin/env python3
from pathlib import Path
import cv2
import depthai as dai
import numpy as np
import time
NNPath = 'YOLOv6n_finetune_V9_best_ckpt_2-5-25.rvc2.tar.xz'
# Create pipeline
with dai.Pipeline() as pipeline:
pipeline.setXLinkChunkSize(0)
# Define the colour camera
ColourCam = pipeline.create(dai.node.Camera).build(dai.CameraBoardSocket.CAM_A, [2024, 1520], 30)
# Resize frame to fit the input of the NN
ColourCamNN = ColourCam.requestOutput([640, 640], dai.ImgFrame.Type.BGR888p, dai.ImgResizeMode.CROP, 30)
# Define the object detector
NNArchive = dai.NNArchive(NNPath)
YOLOv6Network = pipeline.create(dai.node.DetectionNetwork).build(ColourCamNN, NNArchive, 0.8)
YOLOv6Network.setNumNCEPerInferenceThread(1)
YOLOv6Network.setNumInferenceThreads(2)
YOLOv6Network.setNumShavesPerInferenceThread(6)
YOLOv6Network.input.setBlocking(False)
labelMap = YOLOv6Network.getClasses()
# Output queue
qRgb = YOLOv6Network.passthrough.createOutputQueue()
qDet = YOLOv6Network.out.createOutputQueue()
pipeline.start()
frame = None
detections = []
startTime = time.monotonic()
counter = 0
color2 = (255, 255, 255)
# nn data, being the bounding box locations, are in <0..1> range - they need to be normalized with frame width/height
def frameNorm(frame, bbox):
normVals = np.full(len(bbox), frame.shape[0])
normVals[::2] = frame.shape[1]
return (np.clip(np.array(bbox), 0, 1) * normVals).astype(int)
def displayFrame(name, frame):
color = (255, 0, 0)
for detection in detections:
bbox = frameNorm( frame, (detection.xmin, detection.ymin, detection.xmax, detection.ymax),)
cv2.putText(frame, labelMap[detection.label], (bbox[0] + 10, bbox[1] + 20), cv2.FONT_HERSHEY_TRIPLEX, 0.5, 255,)
cv2.putText(frame, f"{int(detection.confidence * 100)}%", (bbox[0] + 10, bbox[1] + 40), cv2.FONT_HERSHEY_TRIPLEX, 0.5, 255,)
cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
# Show the frame
cv2.imshow(name, frame)
while pipeline.isRunning():
#inRgb: dai.ImgFrame = qRgb.get()
#inDet: dai.ImgDetections = qDet.get()
inRgb = qRgb.tryGet()
inDet = qDet.tryGet()
if inRgb is not None:
frame = inRgb.getCvFrame()
cv2.putText(frame, "NN fps: {:.2f}".format(counter / (time.monotonic() - startTime)), (2, frame.shape[0] - 4), cv2.FONT_HERSHEY_TRIPLEX, 0.4, color2,)
if inDet is not None:
detections = inDet.detections
counter += 1
if frame is not None:
displayFrame("rgb", frame)
print("FPS: {:.2f}".format(counter / (time.monotonic() - startTime)))
if cv2.waitKey(1) == ord("q"):
pipeline.stop()
break
