对象追踪器重映射

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    objectTracker = pipeline.create(dai.node.ObjectTracker)
13    labelMap = detectionNetwork.getClasses()
14    depth = pipeline.create(dai.node.Depth).build(dai.node.Depth.Algorithm.AUTO)
15
16    detectionNetwork.out.link(objectTracker.inputDetections)
17    detectionNetwork.passthrough.link(objectTracker.inputDetectionFrame)
18    detectionNetwork.passthrough.link(objectTracker.inputTrackerFrame)
19
20    qRgb = detectionNetwork.passthrough.createOutputQueue()
21    qTrack = objectTracker.out.createOutputQueue()
22    qDepth = depth.depth.createOutputQueue()
23
24    pipeline.start()
25
26    def displayFrame(name: str, frame: dai.ImgFrame, tracklets: dai.Tracklets):
27        color = (0, 255, 0)
28        assert tracklets.getTransformation() is not None
29        if(frame.getType() == dai.ImgFrame.Type.RAW16):
30            cvFrame = dai.utility.colorizeDepthFrame(frame).getCvFrame()
31        else:
32            cvFrame = frame.getCvFrame()
33        for tracklet in tracklets.tracklets:
34            # Get the shape of the frame from which the detections originated for denormalization
35            normShape = tracklets.getTransformation().getSize()
36
37            # Create rotated rectangle to remap
38            # Here we use an intermediate dai.Rect to create a dai.RotatedRect to simplify construction and denormalization
39            rotRect = dai.RotatedRect(tracklet.roi.denormalize(normShape[0], normShape[1]), 0)
40            # Remap the detection rectangle to target frame
41            remapped = tracklets.getTransformation().remapRectTo(frame.getTransformation(), rotRect)
42            # Remapped rectangle could be rotated, so we get the bounding box
43            bbox = [int(l) for l in remapped.getOuterRect()]
44            cv2.putText(
45                cvFrame,
46                labelMap[tracklet.label],
47                (bbox[0] + 10, bbox[1] + 20),
48                cv2.FONT_HERSHEY_TRIPLEX,
49                0.5,
50                255,
51            )
52            cv2.putText(
53                cvFrame,
54                f"{int(tracklet.srcImgDetection.confidence * 100)}%",
55                (bbox[0] + 10, bbox[1] + 40),
56                cv2.FONT_HERSHEY_TRIPLEX,
57                0.5,
58                255,
59            )
60            cv2.rectangle(cvFrame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
61        # Show the frame
62        cv2.imshow(name, cvFrame)
63
64    while pipeline.isRunning():
65        inRgb: dai.ImgFrame = qRgb.get()
66        inTrack: dai.Tracklets = qTrack.get()
67        inDepth: dai.ImgFrame = qDepth.get()
68        hasRgb = inRgb is not None
69        hasDepth = inDepth is not None
70        hasTrack = inTrack is not None
71        if hasRgb:
72            displayFrame("rgb", inRgb, inTrack)
73        if hasDepth:
74            displayFrame("depth", inDepth, inTrack)
75        if cv2.waitKey(1) == ord("q"):
76            pipeline.stop()
77            break

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