jde_tracker.py 16.6 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
15
This code is based on https://github.com/Zhongdao/Towards-Realtime-MOT/blob/master/tracker/multitracker.py
16 17
"""

18 19
import numpy as np
from collections import defaultdict
20 21

from ..matching import jde_matching as matching
22 23
from ..motion import KalmanFilter
from .base_jde_tracker import TrackState, STrack
24 25 26 27 28 29 30 31 32 33 34 35
from .base_jde_tracker import joint_stracks, sub_stracks, remove_duplicate_stracks

from ppdet.core.workspace import register, serializable
from ppdet.utils.logger import setup_logger
logger = setup_logger(__name__)

__all__ = ['JDETracker']


@register
@serializable
class JDETracker(object):
36
    __shared__ = ['num_classes']
37
    """
38
    JDE tracker, support single class and multi classes
39 40

    Args:
41
        use_byte (bool): Whether use ByteTracker, default False
42
        num_classes (int): the number of classes
43 44 45
        det_thresh (float): threshold of detection score
        track_buffer (int): buffer for tracker
        min_box_area (int): min box area to filter out low quality boxes
F
Feng Ni 已提交
46
        vertical_ratio (float): w/h, the vertical ratio of the bbox to filter
47
            bad results. If set <= 0 means no need to filter bboxes,usually set
48
            1.6 for pedestrian tracking.
49 50 51 52 53 54
        tracked_thresh (float): linear assignment threshold of tracked 
            stracks and detections
        r_tracked_thresh (float): linear assignment threshold of 
            tracked stracks and unmatched detections
        unconfirmed_thresh (float): linear assignment threshold of 
            unconfirmed stracks and unmatched detections
55 56 57 58 59 60 61 62
        conf_thres (float): confidence threshold for tracking, also used in
            ByteTracker as higher confidence threshold
        match_thres (float): linear assignment threshold of tracked 
            stracks and detections in ByteTracker
        low_conf_thres (float): lower confidence threshold for tracking in
            ByteTracker
        input_size (list): input feature map size to reid model, [h, w] format,
            [64, 192] as default.
63
        motion (str): motion model, KalmanFilter as default
F
FlyingQianMM 已提交
64 65
        metric_type (str): either "euclidean" or "cosine", the distance metric 
            used for measurement to track association.
66 67 68
    """

    def __init__(self,
F
Feng Ni 已提交
69
                 use_byte=False,
70
                 num_classes=1,
71 72
                 det_thresh=0.3,
                 track_buffer=30,
73 74
                 min_box_area=0,
                 vertical_ratio=0,
75 76 77
                 tracked_thresh=0.7,
                 r_tracked_thresh=0.5,
                 unconfirmed_thresh=0.7,
F
FlyingQianMM 已提交
78
                 conf_thres=0,
F
Feng Ni 已提交
79 80
                 match_thres=0.8,
                 low_conf_thres=0.2,
81
                 input_size=[64, 192],
F
Feng Ni 已提交
82
                 motion='KalmanFilter',
F
FlyingQianMM 已提交
83
                 metric_type='euclidean'):
F
Feng Ni 已提交
84
        self.use_byte = use_byte
85
        self.num_classes = num_classes
F
Feng Ni 已提交
86
        self.det_thresh = det_thresh if not use_byte else conf_thres + 0.1
87 88
        self.track_buffer = track_buffer
        self.min_box_area = min_box_area
F
Feng Ni 已提交
89 90
        self.vertical_ratio = vertical_ratio

91 92 93
        self.tracked_thresh = tracked_thresh
        self.r_tracked_thresh = r_tracked_thresh
        self.unconfirmed_thresh = unconfirmed_thresh
F
Feng Ni 已提交
94 95 96 97
        self.conf_thres = conf_thres
        self.match_thres = match_thres
        self.low_conf_thres = low_conf_thres

98
        self.input_size = input_size
99 100
        if motion == 'KalmanFilter':
            self.motion = KalmanFilter()
F
FlyingQianMM 已提交
101
        self.metric_type = metric_type
102 103

        self.frame_id = 0
104 105 106
        self.tracked_tracks_dict = defaultdict(list)  # dict(list[STrack])
        self.lost_tracks_dict = defaultdict(list)  # dict(list[STrack])
        self.removed_tracks_dict = defaultdict(list)  # dict(list[STrack])
107 108 109 110

        self.max_time_lost = 0
        # max_time_lost will be calculated: int(frame_rate / 30.0 * track_buffer)

F
Feng Ni 已提交
111
    def update(self, pred_dets, pred_embs=None):
112 113 114 115 116 117
        """
        Processes the image frame and finds bounding box(detections).
        Associates the detection with corresponding tracklets and also handles
            lost, removed, refound and active tracklets.

        Args:
118
            pred_dets (np.array): Detection results of the image, the shape is
119
                [N, 6], means 'cls_id, score, x0, y0, x1, y1'.
120 121
            pred_embs (np.array): Embedding results of the image, the shape is
                [N, 128] or [N, 512].
122 123

        Return:
124
            output_stracks_dict (dict(list)): The list contains information
125
                regarding the online_tracklets for the received image tensor.
126 127
        """
        self.frame_id += 1
128 129 130 131 132 133 134
        if self.frame_id == 1:
            STrack.init_count(self.num_classes)
        activated_tracks_dict = defaultdict(list)
        refined_tracks_dict = defaultdict(list)
        lost_tracks_dict = defaultdict(list)
        removed_tracks_dict = defaultdict(list)
        output_tracks_dict = defaultdict(list)
135

136 137
        pred_dets_dict = defaultdict(list)
        pred_embs_dict = defaultdict(list)
F
FlyingQianMM 已提交
138

139 140
        # unify single and multi classes detection and embedding results
        for cls_id in range(self.num_classes):
141
            cls_idx = (pred_dets[:, 0:1] == cls_id).squeeze(-1)
142
            pred_dets_dict[cls_id] = pred_dets[cls_idx]
F
Feng Ni 已提交
143 144 145 146
            if pred_embs is not None:
                pred_embs_dict[cls_id] = pred_embs[cls_idx]
            else:
                pred_embs_dict[cls_id] = None
147

148 149 150 151
        for cls_id in range(self.num_classes):
            """ Step 1: Get detections by class"""
            pred_dets_cls = pred_dets_dict[cls_id]
            pred_embs_cls = pred_embs_dict[cls_id]
152
            remain_inds = (pred_dets_cls[:, 1:2] > self.conf_thres).squeeze(-1)
153 154
            if remain_inds.sum() > 0:
                pred_dets_cls = pred_dets_cls[remain_inds]
155 156
                if pred_embs_cls is None:
                    # in original ByteTrack
F
Feng Ni 已提交
157 158
                    detections = [
                        STrack(
159 160 161 162 163
                            STrack.tlbr_to_tlwh(tlbrs[2:6]),
                            tlbrs[1],
                            cls_id,
                            30,
                            temp_feat=None) for tlbrs in pred_dets_cls
F
Feng Ni 已提交
164 165 166 167 168
                    ]
                else:
                    pred_embs_cls = pred_embs_cls[remain_inds]
                    detections = [
                        STrack(
169
                            STrack.tlbr_to_tlwh(tlbrs[2:6]), tlbrs[1], cls_id,
170 171
                            30, temp_feat) for (tlbrs, temp_feat) in
                        zip(pred_dets_cls, pred_embs_cls)
F
Feng Ni 已提交
172
                    ]
173
            else:
174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189
                detections = []
            ''' Add newly detected tracklets to tracked_stracks'''
            unconfirmed_dict = defaultdict(list)
            tracked_tracks_dict = defaultdict(list)
            for track in self.tracked_tracks_dict[cls_id]:
                if not track.is_activated:
                    # previous tracks which are not active in the current frame are added in unconfirmed list
                    unconfirmed_dict[cls_id].append(track)
                else:
                    # Active tracks are added to the local list 'tracked_stracks'
                    tracked_tracks_dict[cls_id].append(track)
            """ Step 2: First association, with embedding"""
            # building tracking pool for the current frame
            track_pool_dict = defaultdict(list)
            track_pool_dict[cls_id] = joint_stracks(
                tracked_tracks_dict[cls_id], self.lost_tracks_dict[cls_id])
190

191 192
            # Predict the current location with KalmanFilter
            STrack.multi_predict(track_pool_dict[cls_id], self.motion)
193

194 195
            if pred_embs_cls is None:
                # in original ByteTrack
196 197
                dists = matching.iou_distance(track_pool_dict[cls_id],
                                              detections)
F
Feng Ni 已提交
198
                matches, u_track, u_detection = matching.linear_assignment(
199
                    dists, thresh=self.match_thres)  # not self.tracked_thresh
F
Feng Ni 已提交
200 201
            else:
                dists = matching.embedding_distance(
202 203 204 205 206
                    track_pool_dict[cls_id],
                    detections,
                    metric=self.metric_type)
                dists = matching.fuse_motion(
                    self.motion, dists, track_pool_dict[cls_id], detections)
F
Feng Ni 已提交
207 208
                matches, u_track, u_detection = matching.linear_assignment(
                    dists, thresh=self.tracked_thresh)
209

210 211 212 213 214 215 216 217 218 219 220 221 222
            for i_tracked, idet in matches:
                # i_tracked is the id of the track and idet is the detection
                track = track_pool_dict[cls_id][i_tracked]
                det = detections[idet]
                if track.state == TrackState.Tracked:
                    # If the track is active, add the detection to the track
                    track.update(detections[idet], self.frame_id)
                    activated_tracks_dict[cls_id].append(track)
                else:
                    # We have obtained a detection from a track which is not active,
                    # hence put the track in refind_stracks list
                    track.re_activate(det, self.frame_id, new_id=False)
                    refined_tracks_dict[cls_id].append(track)
223

224 225
            # None of the steps below happen if there are no undetected tracks.
            """ Step 3: Second association, with IOU"""
F
Feng Ni 已提交
226
            if self.use_byte:
227 228
                inds_low = pred_dets_dict[cls_id][:, 1:2] > self.low_conf_thres
                inds_high = pred_dets_dict[cls_id][:, 1:2] < self.conf_thres
F
Feng Ni 已提交
229 230
                inds_second = np.logical_and(inds_low, inds_high).squeeze(-1)
                pred_dets_cls_second = pred_dets_dict[cls_id][inds_second]
231

F
Feng Ni 已提交
232 233
                # association the untrack to the low score detections
                if len(pred_dets_cls_second) > 0:
234 235 236 237 238 239 240 241 242 243 244 245
                    if pred_embs_dict[cls_id] is None:
                        # in original ByteTrack
                        detections_second = [
                            STrack(
                                STrack.tlbr_to_tlwh(tlbrs[2:6]),
                                tlbrs[1],
                                cls_id,
                                30,
                                temp_feat=None)
                            for tlbrs in pred_dets_cls_second
                        ]
                    else:
246 247
                        pred_embs_cls_second = pred_embs_dict[cls_id][
                            inds_second]
248 249
                        detections_second = [
                            STrack(
250 251 252
                                STrack.tlbr_to_tlwh(tlbrs[2:6]), tlbrs[1],
                                cls_id, 30, temp_feat) for (tlbrs, temp_feat) in
                            zip(pred_dets_cls_second, pred_embs_cls_second)
253
                        ]
F
Feng Ni 已提交
254 255 256 257 258 259
                else:
                    detections_second = []
                r_tracked_stracks = [
                    track_pool_dict[cls_id][i] for i in u_track
                    if track_pool_dict[cls_id][i].state == TrackState.Tracked
                ]
260 261
                dists = matching.iou_distance(r_tracked_stracks,
                                              detections_second)
F
Feng Ni 已提交
262
                matches, u_track, u_detection_second = matching.linear_assignment(
263
                    dists, thresh=0.4)  # not r_tracked_thresh
F
Feng Ni 已提交
264 265 266 267 268 269 270 271 272 273
            else:
                detections = [detections[i] for i in u_detection]
                r_tracked_stracks = []
                for i in u_track:
                    if track_pool_dict[cls_id][i].state == TrackState.Tracked:
                        r_tracked_stracks.append(track_pool_dict[cls_id][i])
                dists = matching.iou_distance(r_tracked_stracks, detections)

                matches, u_track, u_detection = matching.linear_assignment(
                    dists, thresh=self.r_tracked_thresh)
274

275 276
            for i_tracked, idet in matches:
                track = r_tracked_stracks[i_tracked]
277 278
                det = detections[
                    idet] if not self.use_byte else detections_second[idet]
279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302
                if track.state == TrackState.Tracked:
                    track.update(det, self.frame_id)
                    activated_tracks_dict[cls_id].append(track)
                else:
                    track.re_activate(det, self.frame_id, new_id=False)
                    refined_tracks_dict[cls_id].append(track)

            for it in u_track:
                track = r_tracked_stracks[it]
                if not track.state == TrackState.Lost:
                    track.mark_lost()
                    lost_tracks_dict[cls_id].append(track)
            '''Deal with unconfirmed tracks, usually tracks with only one beginning frame'''
            detections = [detections[i] for i in u_detection]
            dists = matching.iou_distance(unconfirmed_dict[cls_id], detections)
            matches, u_unconfirmed, u_detection = matching.linear_assignment(
                dists, thresh=self.unconfirmed_thresh)
            for i_tracked, idet in matches:
                unconfirmed_dict[cls_id][i_tracked].update(detections[idet],
                                                           self.frame_id)
                activated_tracks_dict[cls_id].append(unconfirmed_dict[cls_id][
                    i_tracked])
            for it in u_unconfirmed:
                track = unconfirmed_dict[cls_id][it]
303
                track.mark_removed()
304 305 306 307 308 309 310 311 312 313 314 315 316
                removed_tracks_dict[cls_id].append(track)
            """ Step 4: Init new stracks"""
            for inew in u_detection:
                track = detections[inew]
                if track.score < self.det_thresh:
                    continue
                track.activate(self.motion, self.frame_id)
                activated_tracks_dict[cls_id].append(track)
            """ Step 5: Update state"""
            for track in self.lost_tracks_dict[cls_id]:
                if self.frame_id - track.end_frame > self.max_time_lost:
                    track.mark_removed()
                    removed_tracks_dict[cls_id].append(track)
317

318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335
            self.tracked_tracks_dict[cls_id] = [
                t for t in self.tracked_tracks_dict[cls_id]
                if t.state == TrackState.Tracked
            ]
            self.tracked_tracks_dict[cls_id] = joint_stracks(
                self.tracked_tracks_dict[cls_id], activated_tracks_dict[cls_id])
            self.tracked_tracks_dict[cls_id] = joint_stracks(
                self.tracked_tracks_dict[cls_id], refined_tracks_dict[cls_id])
            self.lost_tracks_dict[cls_id] = sub_stracks(
                self.lost_tracks_dict[cls_id], self.tracked_tracks_dict[cls_id])
            self.lost_tracks_dict[cls_id].extend(lost_tracks_dict[cls_id])
            self.lost_tracks_dict[cls_id] = sub_stracks(
                self.lost_tracks_dict[cls_id], self.removed_tracks_dict[cls_id])
            self.removed_tracks_dict[cls_id].extend(removed_tracks_dict[cls_id])
            self.tracked_tracks_dict[cls_id], self.lost_tracks_dict[
                cls_id] = remove_duplicate_stracks(
                    self.tracked_tracks_dict[cls_id],
                    self.lost_tracks_dict[cls_id])
336

337 338 339 340 341
            # get scores of lost tracks
            output_tracks_dict[cls_id] = [
                track for track in self.tracked_tracks_dict[cls_id]
                if track.is_activated
            ]
342

343 344 345 346 347 348 349 350 351
            logger.debug('===========Frame {}=========='.format(self.frame_id))
            logger.debug('Activated: {}'.format(
                [track.track_id for track in activated_tracks_dict[cls_id]]))
            logger.debug('Refind: {}'.format(
                [track.track_id for track in refined_tracks_dict[cls_id]]))
            logger.debug('Lost: {}'.format(
                [track.track_id for track in lost_tracks_dict[cls_id]]))
            logger.debug('Removed: {}'.format(
                [track.track_id for track in removed_tracks_dict[cls_id]]))
352

353
        return output_tracks_dict