Object Tracker Replay

1#!/usr/bin/env python3
2
3import cv2
4import depthai as dai
5import time
6
7from pathlib import Path
8from argparse import ArgumentParser
9
10scriptDir = Path(__file__).resolve().parent
11examplesRoot = (scriptDir / Path('../')).resolve()  # This resolves the parent directory correctly
12models = examplesRoot / 'models'
13videoPath = models / 'construction_vest.mp4'
14
15parser = ArgumentParser()
16parser.add_argument("-i", "--inputVideo", default=videoPath, help="Input video name")
17parser.add_argument("-c", "--camera", type=bool, help="Use camera as input", default=False)
18args = parser.parse_args()
19
20# Create pipeline
21with dai.Pipeline() as pipeline:
22    # Define sources and outputs
23    inputSource = None
24    if args.camera: 
25        camRgb = pipeline.create(dai.node.Camera).build(dai.CameraBoardSocket.CAM_A)
26        inputSource = camRgb
27    else:
28        replay = pipeline.create(dai.node.ReplayVideo)
29        replay.setReplayVideoFile(Path(args.inputVideo))
30        inputSource = replay
31
32    detectionNetwork = pipeline.create(dai.node.DetectionNetwork).build(inputSource, "yolov6-nano")
33    objectTracker = pipeline.create(dai.node.ObjectTracker)
34
35    detectionNetwork.setConfidenceThreshold(0.6)
36    detectionNetwork.input.setBlocking(False)
37    labelMap = detectionNetwork.getClasses()
38
39    objectTracker.setDetectionLabelsToTrack([0])  # track only person
40    # possible tracking types: ZERO_TERM_COLOR_HISTOGRAM, ZERO_TERM_IMAGELESS, SHORT_TERM_IMAGELESS, SHORT_TERM_KCF
41    objectTracker.setTrackerType(dai.TrackerType.SHORT_TERM_IMAGELESS)
42    # take the smallest ID when new object is tracked, possible options: SMALLEST_ID, UNIQUE_ID
43    objectTracker.setTrackerIdAssignmentPolicy(dai.TrackerIdAssignmentPolicy.SMALLEST_ID)
44
45    preview = objectTracker.passthroughTrackerFrame.createOutputQueue()
46    tracklets = objectTracker.out.createOutputQueue()
47
48    detectionNetwork.passthrough.link(objectTracker.inputTrackerFrame)
49
50    detectionNetwork.passthrough.link(objectTracker.inputDetectionFrame)
51    detectionNetwork.out.link(objectTracker.inputDetections)
52
53    startTime = time.monotonic()
54    counter = 0
55    fps = 0
56    color = (255, 255, 255)
57    pipeline.start()
58    while(pipeline.isRunning()):
59        imgFrame = preview.get()
60        track = tracklets.get()
61        assert isinstance(imgFrame, dai.ImgFrame), "Expected ImgFrame"
62        assert isinstance(track, dai.Tracklets), "Expected Tracklets"
63
64        counter+=1
65        current_time = time.monotonic()
66        if (current_time - startTime) > 1 :
67            fps = counter / (current_time - startTime)
68            counter = 0
69            startTime = current_time
70
71        frame = imgFrame.getCvFrame()
72        trackletsData = track.tracklets
73        for t in trackletsData:
74            roi = t.roi.denormalize(frame.shape[1], frame.shape[0])
75            x1 = int(roi.topLeft().x)
76            y1 = int(roi.topLeft().y)
77            x2 = int(roi.bottomRight().x)
78            y2 = int(roi.bottomRight().y)
79
80            try:
81                label = labelMap[t.label]
82            except:
83                label = t.label
84
85            cv2.putText(frame, str(label), (x1 + 10, y1 + 20), cv2.FONT_HERSHEY_TRIPLEX, 0.5, 255)
86            cv2.putText(frame, f"ID: {[t.id]}", (x1 + 10, y1 + 35), cv2.FONT_HERSHEY_TRIPLEX, 0.5, 255)
87            cv2.putText(frame, t.status.name, (x1 + 10, y1 + 50), cv2.FONT_HERSHEY_TRIPLEX, 0.5, 255)
88            cv2.rectangle(frame, (x1, y1), (x2, y2), color, cv2.FONT_HERSHEY_SIMPLEX)
89
90        cv2.putText(frame, "NN fps: {:.2f}".format(fps), (2, frame.shape[0] - 4), cv2.FONT_HERSHEY_TRIPLEX, 0.4, color)
91
92        cv2.imshow("tracker", frame)
93
94        if cv2.waitKey(1) == ord('q'):
95            break

1#include <chrono>
2#include <depthai/depthai.hpp>
3#include <opencv2/opencv.hpp>
4
5#include "depthai/pipeline/datatype/Tracklets.hpp"
6
7int main() {
8    // Create pipeline
9    dai::Pipeline pipeline;
10
11    // Define sources and outputs
12    auto replay = pipeline.create<dai::node::ReplayVideo>();
13    replay->setReplayVideoFile(VIDEO_PATH);
14
15    // Create spatial detection network
16    dai::NNModelDescription modelDescription{"yolov6-nano"};
17    auto detectionNetwork = pipeline.create<dai::node::DetectionNetwork>()->build(replay, modelDescription);
18    detectionNetwork->setConfidenceThreshold(0.6f);
19    detectionNetwork->input.setBlocking(false);
20
21    // Create object tracker
22    auto objectTracker = pipeline.create<dai::node::ObjectTracker>();
23    objectTracker->setDetectionLabelsToTrack({0});  // track only person
24    objectTracker->setTrackerType(dai::TrackerType::SHORT_TERM_IMAGELESS);
25    objectTracker->setTrackerIdAssignmentPolicy(dai::TrackerIdAssignmentPolicy::SMALLEST_ID);
26
27    // Create output queues
28    auto preview = objectTracker->passthroughTrackerFrame.createOutputQueue();
29    auto tracklets = objectTracker->out.createOutputQueue();
30
31    // Link nodes
32    detectionNetwork->passthrough.link(objectTracker->inputTrackerFrame);
33
34    detectionNetwork->passthrough.link(objectTracker->inputDetectionFrame);
35    detectionNetwork->out.link(objectTracker->inputDetections);
36
37    // Start pipeline
38    pipeline.start();
39
40    // FPS calculation variables
41    auto startTime = std::chrono::steady_clock::now();
42    int counter = 0;
43    float fps = 0;
44    cv::Scalar color(255, 255, 255);
45
46    while(pipeline.isRunning()) {
47        auto imgFrame = preview->get<dai::ImgFrame>();
48        auto track = tracklets->get<dai::Tracklets>();
49
50        counter++;
51        auto currentTime = std::chrono::steady_clock::now();
52        auto elapsed = std::chrono::duration_cast<std::chrono::seconds>(currentTime - startTime).count();
53        if(elapsed >= 1) {
54            fps = counter / static_cast<float>(elapsed);
55            counter = 0;
56            startTime = currentTime;
57        }
58
59        cv::Mat frame = imgFrame->getCvFrame();
60        auto trackletsData = track->tracklets;
61
62        for(const auto& t : trackletsData) {
63            auto roi = t.roi.denormalize(frame.cols, frame.rows);
64            int x1 = static_cast<int>(roi.topLeft().x);
65            int y1 = static_cast<int>(roi.topLeft().y);
66            int x2 = static_cast<int>(roi.bottomRight().x);
67            int y2 = static_cast<int>(roi.bottomRight().y);
68
69            std::string label;
70            try {
71                label = detectionNetwork->getClasses().value()[t.label];
72            } catch(...) {
73                label = std::to_string(t.label);
74            }
75
76            cv::putText(frame, label, cv::Point(x1 + 10, y1 + 20), cv::FONT_HERSHEY_TRIPLEX, 0.5, color);
77            cv::putText(frame, "ID: " + std::to_string(t.id), cv::Point(x1 + 10, y1 + 35), cv::FONT_HERSHEY_TRIPLEX, 0.5, color);
78            cv::putText(frame,
79                        std::string(t.status == dai::Tracklet::TrackingStatus::LOST ? "LOST" : "TRACKED"),
80                        cv::Point(x1 + 10, y1 + 50),
81                        cv::FONT_HERSHEY_TRIPLEX,
82                        0.5,
83                        color);
84            cv::rectangle(frame, cv::Point(x1, y1), cv::Point(x2, y2), color, cv::FONT_HERSHEY_SIMPLEX);
85        }
86
87        cv::putText(frame, "NN fps: " + std::to_string(fps).substr(0, 4), cv::Point(2, frame.rows - 4), cv::FONT_HERSHEY_TRIPLEX, 0.4, color);
88
89        cv::imshow("tracker", frame);
90
91        if(cv::waitKey(1) == 'q') {
92            break;
93        }
94    }
95
96    return 0;
97}