检测网络重映射

1#!/usr/bin/env python3
2
3import cv2
4import depthai as dai
5
6# Create pipeline
7with dai.Pipeline() as pipeline:
8    colorSockets = pipeline.getDefaultDevice().getConnectedCameras(dai.CameraSensorType.COLOR)
9    colorSocket = colorSockets[0] if colorSockets else dai.CameraBoardSocket.CAM_A
10    cameraNode = pipeline.create(dai.node.Camera).build(colorSocket)
11    detectionNetwork = pipeline.create(dai.node.DetectionNetwork).build(cameraNode, dai.NNModelDescription("yolov6-nano"))
12    labelMap = detectionNetwork.getClasses()
13    depth = pipeline.create(dai.node.Depth).build(dai.node.Depth.Algorithm.AUTO, None)
14
15    qRgb = detectionNetwork.passthrough.createOutputQueue()
16    qDet = detectionNetwork.out.createOutputQueue()
17    qDepth = depth.depth.createOutputQueue()
18
19    pipeline.start()
20
21    def displayFrame(name: str, frame: dai.ImgFrame, imgDetections: dai.ImgDetections):
22        color = (0, 255, 0)
23        assert imgDetections.getTransformation() is not None
24        if(frame.getType() == dai.ImgFrame.Type.RAW16):
25            cvFrame = dai.utility.colorizeDepthFrame(frame).getCvFrame()
26        else:
27            cvFrame = frame.getCvFrame()
28        for detection in imgDetections.detections:
29            # Get the shape of the frame from which the detections originated for denormalization
30            normShape = imgDetections.getTransformation().getSize()
31
32            # Create rotated rectangle to remap
33            # Here we use an intermediate dai.Rect to create a dai.RotatedRect to simplify construction and denormalization
34            rotRect = dai.RotatedRect(dai.Rect(dai.Point2f(detection.xmin, detection.ymin), dai.Point2f(detection.xmax, detection.ymax)).denormalize(normShape[0], normShape[1]), 0)
35            # Remap the detection rectangle to target frame
36            remapped = imgDetections.getTransformation().remapRectTo(frame.getTransformation(), rotRect)
37            # Remapped rectangle could be rotated, so we get the bounding box
38            bbox = [int(l) for l in remapped.getOuterRect()]
39            cv2.putText(
40                cvFrame,
41                labelMap[detection.label],
42                (bbox[0] + 10, bbox[1] + 20),
43                cv2.FONT_HERSHEY_TRIPLEX,
44                0.5,
45                255,
46            )
47            cv2.putText(
48                cvFrame,
49                f"{int(detection.confidence * 100)}%",
50                (bbox[0] + 10, bbox[1] + 40),
51                cv2.FONT_HERSHEY_TRIPLEX,
52                0.5,
53                255,
54            )
55            cv2.rectangle(cvFrame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
56        # Show the frame
57        cv2.imshow(name, cvFrame)
58
59    while pipeline.isRunning():
60        inRgb: dai.ImgFrame = qRgb.get()
61        inDet: dai.ImgDetections = qDet.get()
62        inDepth: dai.ImgFrame = qDepth.get()
63        hasRgb = inRgb is not None
64        hasDepth = inDepth is not None
65        hasDet = inDet is not None
66        if hasRgb:
67            displayFrame("rgb", inRgb, inDet)
68        if hasDepth:
69            displayFrame("depth", inDepth, inDet)
70        if cv2.waitKey(1) == ord("q"):
71            pipeline.stop()
72            break

1#include <csignal>
2#include <iostream>
3#include <opencv2/opencv.hpp>
4#include <string>
5#include <vector>
6
7#include "depthai/depthai.hpp"
8
9std::atomic<bool> quitEvent(false);
10
11void signalHandler(int) {
12    quitEvent = true;
13}
14
15// Helper function to display frames with detections
16void displayFrame(const std::string& name,
17                  std::shared_ptr<dai::ImgFrame> frame,
18                  std::shared_ptr<dai::ImgDetections> imgDetections,
19                  const std::vector<std::string>& labelMap) {
20    cv::Scalar color(0, 255, 0);
21    cv::Mat cvFrame;
22
23    if(frame->getType() == dai::ImgFrame::Type::RAW16) {
24        cvFrame = dai::utility::colorizeDepthFrame(*frame).getCvFrame();
25    } else {
26        cvFrame = frame->getCvFrame();
27    }
28
29    if(!imgDetections || !imgDetections->transformation.has_value()) {
30        // std::cout << "No detections or transformation data for " << name << std::endl;
31        cv::imshow(name, cvFrame);
32        return;
33    }
34
35    const auto& sourceTransform = *(imgDetections->transformation);
36    const auto& targetTransform = frame->transformation;
37
38    for(const auto& detection : imgDetections->detections) {
39        auto normShape = sourceTransform.getSize();
40
41        dai::Rect rect(dai::Point2f(detection.xmin, detection.ymin), dai::Point2f(detection.xmax, detection.ymax));
42        rect = rect.denormalize(static_cast<float>(normShape.first), static_cast<float>(normShape.second));
43        dai::RotatedRect rotRect(rect, 0);
44
45        auto remapped = sourceTransform.remapRectTo(targetTransform, rotRect);
46        auto bbox = remapped.getOuterRect();
47
48        cv::putText(cvFrame,
49                    labelMap[detection.label],
50                    cv::Point(static_cast<int>(bbox[0]) + 10, static_cast<int>(bbox[1]) + 20),
51                    cv::FONT_HERSHEY_TRIPLEX,
52                    0.5,
53                    cv::Scalar(255, 255, 255));
54        cv::putText(cvFrame,
55                    std::to_string(static_cast<int>(detection.confidence * 100)) + "%",
56                    cv::Point(static_cast<int>(bbox[0]) + 10, static_cast<int>(bbox[1]) + 40),
57                    cv::FONT_HERSHEY_TRIPLEX,
58                    0.5,
59                    cv::Scalar(255, 255, 255));
60        cv::rectangle(cvFrame,
61                      cv::Point(static_cast<int>(bbox[0]), static_cast<int>(bbox[1])),
62                      cv::Point(static_cast<int>(bbox[2]), static_cast<int>(bbox[3])),
63                      color,
64                      2);
65    }
66    cv::imshow(name, cvFrame);
67}
68
69int main() {
70    signal(SIGTERM, signalHandler);
71    signal(SIGINT, signalHandler);
72
73    dai::Pipeline pipeline;
74
75    auto colorSockets = pipeline.getDefaultDevice()->getConnectedCameras(dai::CameraSensorType::COLOR);
76    auto colorSocket = colorSockets.empty() ? dai::CameraBoardSocket::CAM_A : colorSockets.front();
77    auto cameraNode = pipeline.create<dai::node::Camera>();
78    cameraNode->build(colorSocket);
79
80    auto detectionNetwork = pipeline.create<dai::node::DetectionNetwork>();
81    dai::NNModelDescription modelDescription;
82    modelDescription.model = "yolov6-nano";
83    detectionNetwork->build(cameraNode, modelDescription);
84    auto labelMap = detectionNetwork->getClasses().value_or(std::vector<std::string>{});
85
86    auto depth = pipeline.create<dai::node::Depth>();
87    depth->build(dai::node::Depth::Algorithm::AUTO);
88
89    auto qRgb = detectionNetwork->passthrough.createOutputQueue();
90    auto qDet = detectionNetwork->out.createOutputQueue();
91    auto qDepth = depth->depth().createOutputQueue();
92
93    pipeline.start();
94
95    while(pipeline.isRunning() && !quitEvent) {
96        auto inRgb = qRgb->tryGet<dai::ImgFrame>();
97        auto inDet = qDet->tryGet<dai::ImgDetections>();
98        auto inDepth = qDepth->tryGet<dai::ImgFrame>();
99
100        bool hasRgb = inRgb != nullptr;
101        bool hasDepth = inDepth != nullptr;
102        bool hasDet = inDet != nullptr;
103
104        if(hasRgb && hasDet) {
105            displayFrame("rgb", inRgb, inDet, labelMap);
106        }
107        if(hasDepth && hasDet) {
108            displayFrame("depth", inDepth, inDet, labelMap);
109        }
110
111        if(cv::waitKey(1) == 'q') {
112            pipeline.stop();
113            break;
114        }
115    }
116
117    return 0;
118}