Object tracker on RGB

1#!/usr/bin/env python3
2
3from pathlib import Path
4import cv2
5import depthai as dai
6import numpy as np
7import time
8import argparse
9
10labelMap = ["background", "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow",
11            "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
12
13nnPathDefault = str((Path(__file__).parent / Path('../models/mobilenet-ssd_openvino_2021.4_6shave.blob')).resolve().absolute())
14parser = argparse.ArgumentParser()
15parser.add_argument('nnPath', nargs='?', help="Path to mobilenet detection network blob", default=nnPathDefault)
16parser.add_argument('-ff', '--full_frame', action="store_true", help="Perform tracking on full RGB frame", default=False)
17
18args = parser.parse_args()
19
20fullFrameTracking = args.full_frame
21
22# Create pipeline
23pipeline = dai.Pipeline()
24
25# Define sources and outputs
26camRgb = pipeline.create(dai.node.ColorCamera)
27detectionNetwork = pipeline.create(dai.node.MobileNetDetectionNetwork)
28objectTracker = pipeline.create(dai.node.ObjectTracker)
29
30xlinkOut = pipeline.create(dai.node.XLinkOut)
31trackerOut = pipeline.create(dai.node.XLinkOut)
32xinTrackerConfig = pipeline.create(dai.node.XLinkIn)
33
34xlinkOut.setStreamName("preview")
35trackerOut.setStreamName("tracklets")
36xinTrackerConfig.setStreamName("trackerConfig")
37
38# Properties
39camRgb.setPreviewSize(300, 300)
40camRgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
41camRgb.setInterleaved(False)
42camRgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)
43camRgb.setFps(40)
44
45# testing MobileNet DetectionNetwork
46detectionNetwork.setBlobPath(args.nnPath)
47detectionNetwork.setConfidenceThreshold(0.5)
48detectionNetwork.input.setBlocking(False)
49
50objectTracker.setDetectionLabelsToTrack([15])  # track only person
51# possible tracking types: ZERO_TERM_COLOR_HISTOGRAM, ZERO_TERM_IMAGELESS, SHORT_TERM_IMAGELESS, SHORT_TERM_KCF
52objectTracker.setTrackerType(dai.TrackerType.ZERO_TERM_COLOR_HISTOGRAM)
53# take the smallest ID when new object is tracked, possible options: SMALLEST_ID, UNIQUE_ID
54objectTracker.setTrackerIdAssignmentPolicy(dai.TrackerIdAssignmentPolicy.SMALLEST_ID)
55
56# Linking
57camRgb.preview.link(detectionNetwork.input)
58objectTracker.passthroughTrackerFrame.link(xlinkOut.input)
59
60if fullFrameTracking:
61    camRgb.video.link(objectTracker.inputTrackerFrame)
62else:
63    detectionNetwork.passthrough.link(objectTracker.inputTrackerFrame)
64
65detectionNetwork.passthrough.link(objectTracker.inputDetectionFrame)
66detectionNetwork.out.link(objectTracker.inputDetections)
67objectTracker.out.link(trackerOut.input)
68
69# set tracking parameters
70objectTracker.setOcclusionRatioThreshold(0.4)
71objectTracker.setTrackletMaxLifespan(120)
72objectTracker.setTrackletBirthThreshold(3)
73
74xinTrackerConfig.out.link(objectTracker.inputConfig)
75
76# Connect to device and start pipeline
77with dai.Device(pipeline) as device:
78
79    preview = device.getOutputQueue("preview", 4, False)
80    tracklets = device.getOutputQueue("tracklets", 4, False)
81    trackerConfigQueue = device.getInputQueue("trackerConfig")
82
83    startTime = time.monotonic()
84    counter = 0
85    fps = 0
86    frame = None
87
88    while(True):
89        latestTrackedIds = []
90        imgFrame = preview.get()
91        track = tracklets.get()
92
93        counter+=1
94        current_time = time.monotonic()
95        if (current_time - startTime) > 1 :
96            fps = counter / (current_time - startTime)
97            counter = 0
98            startTime = current_time
99
100        color = (255, 0, 0)
101        frame = imgFrame.getCvFrame()
102        trackletsData = track.tracklets
103        for t in trackletsData:
104            roi = t.roi.denormalize(frame.shape[1], frame.shape[0])
105            x1 = int(roi.topLeft().x)
106            y1 = int(roi.topLeft().y)
107            x2 = int(roi.bottomRight().x)
108            y2 = int(roi.bottomRight().y)
109
110            try:
111                label = labelMap[t.label]
112            except:
113                label = t.label
114
115            cv2.putText(frame, str(label), (x1 + 10, y1 + 20), cv2.FONT_HERSHEY_TRIPLEX, 0.5, 255)
116            cv2.putText(frame, f"ID: {[t.id]}", (x1 + 10, y1 + 35), cv2.FONT_HERSHEY_TRIPLEX, 0.5, 255)
117            cv2.putText(frame, t.status.name, (x1 + 10, y1 + 50), cv2.FONT_HERSHEY_TRIPLEX, 0.5, 255)
118            cv2.rectangle(frame, (x1, y1), (x2, y2), color, cv2.FONT_HERSHEY_SIMPLEX)
119
120            if t.status == dai.Tracklet.TrackingStatus.TRACKED:
121                latestTrackedIds.append(t.id)
122
123        cv2.putText(frame, "NN fps: {:.2f}".format(fps), (2, frame.shape[0] - 4), cv2.FONT_HERSHEY_TRIPLEX, 0.4, color)
124
125        cv2.imshow("tracker", frame)
126
127        key = cv2.waitKey(1)
128        if key == ord('q'):
129            break
130        elif key == ord('g'):
131            # send tracker config to device
132            config = dai.ObjectTrackerConfig()
133
134            # take a random ID from the latest tracked IDs
135            if len(latestTrackedIds) > 0:
136                idToRemove = (np.random.choice(latestTrackedIds))
137                print(f"Force removing ID: {idToRemove}")
138                config.forceRemoveID(idToRemove)
139                trackerConfigQueue.send(config)
140            else:
141                print("No tracked IDs available to force remove")

1#include <chrono>
2#include <iostream>
3
4// Includes common necessary includes for development using depthai library
5#include "depthai/depthai.hpp"
6
7static const std::vector<std::string> labelMap = {"background", "aeroplane", "bicycle",     "bird",  "boat",        "bottle", "bus",
8                                                  "car",        "cat",       "chair",       "cow",   "diningtable", "dog",    "horse",
9                                                  "motorbike",  "person",    "pottedplant", "sheep", "sofa",        "train",  "tvmonitor"};
10
11static std::atomic<bool> fullFrameTracking{false};
12
13int main(int argc, char** argv) {
14    using namespace std;
15    using namespace std::chrono;
16    std::string nnPath(BLOB_PATH);
17
18    // If path to blob specified, use that
19    if(argc > 1) {
20        nnPath = std::string(argv[1]);
21    }
22
23    // Print which blob we are using
24    printf("Using blob at path: %s\n", nnPath.c_str());
25
26    // Create pipeline
27    dai::Pipeline pipeline;
28
29    // Define sources and outputs
30    auto camRgb = pipeline.create<dai::node::ColorCamera>();
31    auto detectionNetwork = pipeline.create<dai::node::MobileNetDetectionNetwork>();
32    auto objectTracker = pipeline.create<dai::node::ObjectTracker>();
33
34    auto xlinkOut = pipeline.create<dai::node::XLinkOut>();
35    auto trackerOut = pipeline.create<dai::node::XLinkOut>();
36
37    xlinkOut->setStreamName("preview");
38    trackerOut->setStreamName("tracklets");
39
40    // Properties
41    camRgb->setPreviewSize(300, 300);
42    camRgb->setResolution(dai::ColorCameraProperties::SensorResolution::THE_1080_P);
43    camRgb->setInterleaved(false);
44    camRgb->setColorOrder(dai::ColorCameraProperties::ColorOrder::BGR);
45    camRgb->setFps(40);
46
47    // testing MobileNet DetectionNetwork
48    detectionNetwork->setBlobPath(nnPath);
49    detectionNetwork->setConfidenceThreshold(0.5f);
50    detectionNetwork->input.setBlocking(false);
51
52    objectTracker->setDetectionLabelsToTrack({15});  // track only person
53    // possible tracking types: ZERO_TERM_COLOR_HISTOGRAM, ZERO_TERM_IMAGELESS, SHORT_TERM_IMAGELESS, SHORT_TERM_KCF
54    objectTracker->setTrackerType(dai::TrackerType::ZERO_TERM_COLOR_HISTOGRAM);
55    // take the smallest ID when new object is tracked, possible options: SMALLEST_ID, UNIQUE_ID
56    objectTracker->setTrackerIdAssignmentPolicy(dai::TrackerIdAssignmentPolicy::SMALLEST_ID);
57
58    // Linking
59    camRgb->preview.link(detectionNetwork->input);
60    objectTracker->passthroughTrackerFrame.link(xlinkOut->input);
61
62    if(fullFrameTracking) {
63        camRgb->video.link(objectTracker->inputTrackerFrame);
64    } else {
65        detectionNetwork->passthrough.link(objectTracker->inputTrackerFrame);
66    }
67
68    detectionNetwork->passthrough.link(objectTracker->inputDetectionFrame);
69    detectionNetwork->out.link(objectTracker->inputDetections);
70    objectTracker->out.link(trackerOut->input);
71
72    // Connect to device and start pipeline
73    dai::Device device(pipeline);
74
75    auto preview = device.getOutputQueue("preview", 4, false);
76    auto tracklets = device.getOutputQueue("tracklets", 4, false);
77
78    auto startTime = steady_clock::now();
79    int counter = 0;
80    float fps = 0;
81
82    while(true) {
83        auto imgFrame = preview->get<dai::ImgFrame>();
84        auto track = tracklets->get<dai::Tracklets>();
85
86        counter++;
87        auto currentTime = steady_clock::now();
88        auto elapsed = duration_cast<duration<float>>(currentTime - startTime);
89        if(elapsed > seconds(1)) {
90            fps = counter / elapsed.count();
91            counter = 0;
92            startTime = currentTime;
93        }
94
95        auto color = cv::Scalar(255, 0, 0);
96        cv::Mat frame = imgFrame->getCvFrame();
97        auto trackletsData = track->tracklets;
98        for(auto& t : trackletsData) {
99            auto roi = t.roi.denormalize(frame.cols, frame.rows);
100            int x1 = roi.topLeft().x;
101            int y1 = roi.topLeft().y;
102            int x2 = roi.bottomRight().x;
103            int y2 = roi.bottomRight().y;
104
105            uint32_t labelIndex = t.label;
106            std::string labelStr = to_string(labelIndex);
107            if(labelIndex < labelMap.size()) {
108                labelStr = labelMap[labelIndex];
109            }
110            cv::putText(frame, labelStr, cv::Point(x1 + 10, y1 + 20), cv::FONT_HERSHEY_TRIPLEX, 0.5, color);
111
112            std::stringstream idStr;
113            idStr << "ID: " << t.id;
114            cv::putText(frame, idStr.str(), cv::Point(x1 + 10, y1 + 40), cv::FONT_HERSHEY_TRIPLEX, 0.5, color);
115            std::stringstream statusStr;
116            statusStr << "Status: " << t.status;
117            cv::putText(frame, statusStr.str(), cv::Point(x1 + 10, y1 + 60), cv::FONT_HERSHEY_TRIPLEX, 0.5, color);
118
119            cv::rectangle(frame, cv::Rect(cv::Point(x1, y1), cv::Point(x2, y2)), color, cv::FONT_HERSHEY_SIMPLEX);
120        }
121
122        std::stringstream fpsStr;
123        fpsStr << "NN fps:" << std::fixed << std::setprecision(2) << fps;
124        cv::putText(frame, fpsStr.str(), cv::Point(2, imgFrame->getHeight() - 4), cv::FONT_HERSHEY_TRIPLEX, 0.4, color);
125
126        cv::imshow("tracker", frame);
127
128        int key = cv::waitKey(1);
129        if(key == 'q' || key == 'Q') {
130            return 0;
131        }
132    }
133    return 0;
134}