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