AprilTags on image replay

View as Markdown

1#!/usr/bin/env python3
2
3import cv2
4import depthai as dai
5import time
6from pathlib import Path
7
8# Get the absolute path of the current script's directory
9script_dir = Path(__file__).resolve().parent
10examplesRoot = (script_dir / Path('../')).resolve()  # This resolves the parent directory correctly
11models = examplesRoot / 'models'
12tagImage = models / 'april_tags.jpg'
13
14class ImageReplay(dai.node.ThreadedHostNode):
15    def __init__(self):
16        dai.node.ThreadedHostNode.__init__(self)
17        self.output =  self.createOutput()
18        frame = cv2.imread(str(tagImage.resolve()))
19        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
20        imgFrame = dai.ImgFrame()
21        imgFrame.setData(frame)
22        imgFrame.setWidth(frame.shape[1])
23        imgFrame.setHeight(frame.shape[0])
24        imgFrame.setType(dai.ImgFrame.Type.GRAY8)
25        self.imgFrame = imgFrame
26    def run(self):
27        while self.mainLoop():
28            self.output.send(self.imgFrame)
29            time.sleep(0.03)
30
31with dai.Pipeline() as pipeline:
32    imageReplay = pipeline.create(ImageReplay)
33    aprilTagNode = pipeline.create(dai.node.AprilTag)
34    imageReplay.output.link(aprilTagNode.inputImage)
35    aprilTagNode.initialConfig.setFamily(dai.AprilTagConfig.Family.TAG_16H5)
36
37    passthroughOutputQueue = aprilTagNode.passthroughInputImage.createOutputQueue()
38    outQueue = aprilTagNode.out.createOutputQueue()
39
40    color = (0, 255, 0)
41    startTime = time.monotonic()
42    counter = 0
43    fps = 0.0
44
45    pipeline.start()
46    while pipeline.isRunning():
47        aprilTagMessage = outQueue.get()
48        assert(isinstance(aprilTagMessage, dai.AprilTags))
49        aprilTags = aprilTagMessage.aprilTags
50
51        counter += 1
52        currentTime = time.monotonic()
53        if (currentTime - startTime) > 1:
54            fps = counter / (currentTime - startTime)
55            counter = 0
56            startTime = currentTime
57
58        passthroughImage: dai.ImgFrame = passthroughOutputQueue.get()
59        frame = passthroughImage.getCvFrame()
60        frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
61
62        def to_int(tag):
63            return (int(tag.x), int(tag.y))
64
65        for tag in aprilTags:
66            topLeft = to_int(tag.topLeft)
67            topRight = to_int(tag.topRight)
68            bottomRight = to_int(tag.bottomRight)
69            bottomLeft = to_int(tag.bottomLeft)
70
71            center = (int((topLeft[0] + bottomRight[0]) / 2), int((topLeft[1] + bottomRight[1]) / 2))
72
73            cv2.line(frame, topLeft, topRight, color, 2, cv2.LINE_AA, 0)
74            cv2.line(frame, topRight,bottomRight, color, 2, cv2.LINE_AA, 0)
75            cv2.line(frame, bottomRight,bottomLeft, color, 2, cv2.LINE_AA, 0)
76            cv2.line(frame, bottomLeft,topLeft, color, 2, cv2.LINE_AA, 0)
77
78            idStr = "ID: " + str(tag.id)
79            cv2.putText(frame, idStr, center, cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
80
81            cv2.putText(frame, f"fps: {fps:.1f}", (200, 20), cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
82
83        cv2.imshow("detections", frame)
84        if cv2.waitKey(1) == ord("q"):
85            break

1#include <atomic>
2#include <chrono>
3#include <csignal>
4#include <iostream>
5#include <opencv2/opencv.hpp>
6#include <thread>
7
8#include "depthai/depthai.hpp"
9
10std::atomic<bool> quitEvent(false);
11
12void signalHandler(int) {
13    quitEvent = true;
14}
15
16class ImageReplay : public dai::NodeCRTP<dai::node::ThreadedHostNode, ImageReplay> {
17   public:
18    constexpr static const char* NAME = "ImageReplay";
19
20    Output output{*this, {"out", DEFAULT_GROUP, {{{dai::DatatypeEnum::ImgFrame, true}}}}};
21
22    ImageReplay() {
23        // Load and prepare the image
24        cv::Mat frame = cv::imread(APRIL_TAGS_PATH);
25        cv::cvtColor(frame, frame, cv::COLOR_BGR2GRAY);
26
27        // Create ImgFrame
28        auto imgFrame = std::make_shared<dai::ImgFrame>();
29        std::vector<uint8_t> data(frame.data, frame.data + frame.total() * frame.elemSize());
30        imgFrame->setData(data);
31        imgFrame->setWidth(frame.cols);
32        imgFrame->setHeight(frame.rows);
33        imgFrame->setType(dai::ImgFrame::Type::GRAY8);
34        _imgFrame = imgFrame;
35    }
36
37    void run() override {
38        while(mainLoop()) {
39            output.send(_imgFrame);
40            std::this_thread::sleep_for(std::chrono::milliseconds(30));
41        }
42    }
43
44   private:
45    std::shared_ptr<dai::ImgFrame> _imgFrame;
46};
47
48int main() {
49    signal(SIGTERM, signalHandler);
50    signal(SIGINT, signalHandler);
51
52    // Create pipeline
53    dai::Pipeline pipeline;
54
55    // Create nodes
56    auto imageReplay = pipeline.create<ImageReplay>();
57    auto aprilTagNode = pipeline.create<dai::node::AprilTag>();
58
59    // Link nodes
60    imageReplay->output.link(aprilTagNode->inputImage);
61    aprilTagNode->initialConfig->setFamily(dai::AprilTagConfig::Family::TAG_16H5);
62
63    // Create output queues
64    auto passthroughOutputQueue = aprilTagNode->passthroughInputImage.createOutputQueue();
65    auto outQueue = aprilTagNode->out.createOutputQueue();
66
67    // Start pipeline
68    pipeline.start();
69
70    // FPS calculation variables
71    cv::Scalar color(0, 255, 0);
72    auto startTime = std::chrono::steady_clock::now();
73    int counter = 0;
74    float fps = 0.0f;
75
76    // Main loop
77    while(pipeline.isRunning() && !quitEvent) {
78        auto aprilTagMessage = outQueue->get<dai::AprilTags>();
79        auto aprilTags = aprilTagMessage->aprilTags;
80
81        // Calculate FPS
82        counter++;
83        auto currentTime = std::chrono::steady_clock::now();
84        auto elapsed = std::chrono::duration_cast<std::chrono::seconds>(currentTime - startTime).count();
85        if(elapsed >= 1) {
86            fps = counter / static_cast<float>(elapsed);
87            counter = 0;
88            startTime = currentTime;
89        }
90
91        // Get passthrough image
92        auto passthroughImage = passthroughOutputQueue->get<dai::ImgFrame>();
93        cv::Mat frame = passthroughImage->getCvFrame();
94        cv::cvtColor(frame, frame, cv::COLOR_GRAY2BGR);
95
96        // Helper function to convert points to integers
97        auto to_int = [](const dai::Point2f& p) { return cv::Point(static_cast<int>(p.x), static_cast<int>(p.y)); };
98
99        // Draw detections
100        for(const auto& tag : aprilTags) {
101            auto topLeft = to_int(tag.topLeft);
102            auto topRight = to_int(tag.topRight);
103            auto bottomRight = to_int(tag.bottomRight);
104            auto bottomLeft = to_int(tag.bottomLeft);
105
106            auto center = cv::Point((topLeft.x + bottomRight.x) / 2, (topLeft.y + bottomRight.y) / 2);
107
108            // Draw rectangle
109            cv::line(frame, topLeft, topRight, color, 2, cv::LINE_AA, 0);
110            cv::line(frame, topRight, bottomRight, color, 2, cv::LINE_AA, 0);
111            cv::line(frame, bottomRight, bottomLeft, color, 2, cv::LINE_AA, 0);
112            cv::line(frame, bottomLeft, topLeft, color, 2, cv::LINE_AA, 0);
113
114            // Draw ID
115            std::string idStr = "ID: " + std::to_string(tag.id);
116            cv::putText(frame, idStr, center, cv::FONT_HERSHEY_TRIPLEX, 0.5, color);
117
118            // Draw FPS
119            cv::putText(frame, "fps: " + std::to_string(fps).substr(0, 4), cv::Point(200, 20), cv::FONT_HERSHEY_TRIPLEX, 0.5, color);
120        }
121
122        cv::imshow("detections", frame);
123        if(cv::waitKey(1) == 'q') {
124            break;
125        }
126    }
127
128    pipeline.stop();
129    pipeline.wait();
130
131    return 0;
132}