# 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. from __future__ import absolute_import from __future__ import division from __future__ import print_function import os import cv2 import glob import paddle import numpy as np from ppdet.core.workspace import create from ppdet.utils.checkpoint import load_weight, load_pretrain_weight from ppdet.modeling.mot.utils import Detection, get_crops, scale_coords, clip_box from ppdet.modeling.mot.utils import Timer, load_det_results from ppdet.modeling.mot import visualization as mot_vis from ppdet.metrics import Metric, MOTMetric, KITTIMOTMetric import ppdet.utils.stats as stats from .callbacks import Callback, ComposeCallback from ppdet.utils.logger import setup_logger logger = setup_logger(__name__) __all__ = ['Tracker'] class Tracker(object): def __init__(self, cfg, mode='eval'): self.cfg = cfg assert mode.lower() in ['test', 'eval'], \ "mode should be 'test' or 'eval'" self.mode = mode.lower() self.optimizer = None # build MOT data loader self.dataset = cfg['{}MOTDataset'.format(self.mode.capitalize())] # build model self.model = create(cfg.architecture) self.status = {} self.start_epoch = 0 # initial default callbacks self._init_callbacks() # initial default metrics self._init_metrics() self._reset_metrics() def _init_callbacks(self): self._callbacks = [] self._compose_callback = None def _init_metrics(self): if self.mode in ['test']: self._metrics = [] return if self.cfg.metric == 'MOT': self._metrics = [MOTMetric(), ] elif self.cfg.metric == 'KITTI': self._metrics = [KITTIMOTMetric(), ] else: logger.warning("Metric not support for metric type {}".format( self.cfg.metric)) self._metrics = [] def _reset_metrics(self): for metric in self._metrics: metric.reset() def register_callbacks(self, callbacks): callbacks = [h for h in list(callbacks) if h is not None] for c in callbacks: assert isinstance(c, Callback), \ "metrics shoule be instances of subclass of Metric" self._callbacks.extend(callbacks) self._compose_callback = ComposeCallback(self._callbacks) def register_metrics(self, metrics): metrics = [m for m in list(metrics) if m is not None] for m in metrics: assert isinstance(m, Metric), \ "metrics shoule be instances of subclass of Metric" self._metrics.extend(metrics) def load_weights_jde(self, weights): load_weight(self.model, weights, self.optimizer) def load_weights_sde(self, det_weights, reid_weights): if self.model.detector: load_weight(self.model.detector, det_weights) load_weight(self.model.reid, reid_weights) else: load_weight(self.model.reid, reid_weights, self.optimizer) def _eval_seq_jde(self, dataloader, save_dir=None, show_image=False, frame_rate=30, draw_threshold=0): if save_dir: if not os.path.exists(save_dir): os.makedirs(save_dir) tracker = self.model.tracker tracker.max_time_lost = int(frame_rate / 30.0 * tracker.track_buffer) timer = Timer() results = [] frame_id = 0 self.status['mode'] = 'track' self.model.eval() for step_id, data in enumerate(dataloader): self.status['step_id'] = step_id if frame_id % 40 == 0: logger.info('Processing frame {} ({:.2f} fps)'.format( frame_id, 1. / max(1e-5, timer.average_time))) # forward timer.tic() pred_dets, pred_embs = self.model(data) online_targets = self.model.tracker.update(pred_dets, pred_embs) online_tlwhs, online_ids = [], [] online_scores = [] for t in online_targets: tlwh = t.tlwh tid = t.track_id tscore = t.score if tscore < draw_threshold: continue if tlwh[2] * tlwh[3] <= tracker.min_box_area: continue if tracker.vertical_ratio > 0 and tlwh[2] / tlwh[ 3] > tracker.vertical_ratio: continue online_tlwhs.append(tlwh) online_ids.append(tid) online_scores.append(tscore) timer.toc() # save results results.append( (frame_id + 1, online_tlwhs, online_scores, online_ids)) self.save_results(data, frame_id, online_ids, online_tlwhs, online_scores, timer.average_time, show_image, save_dir) frame_id += 1 return results, frame_id, timer.average_time, timer.calls def _eval_seq_sde(self, dataloader, save_dir=None, show_image=False, frame_rate=30, det_file='', draw_threshold=0): if save_dir: if not os.path.exists(save_dir): os.makedirs(save_dir) tracker = self.model.tracker use_detector = False if not self.model.detector else True timer = Timer() results = [] frame_id = 0 self.status['mode'] = 'track' self.model.eval() self.model.reid.eval() if not use_detector: dets_list = load_det_results(det_file, len(dataloader)) logger.info('Finish loading detection results file {}.'.format( det_file)) for step_id, data in enumerate(dataloader): self.status['step_id'] = step_id if frame_id % 40 == 0: logger.info('Processing frame {} ({:.2f} fps)'.format( frame_id, 1. / max(1e-5, timer.average_time))) ori_image = data['ori_image'] input_shape = data['image'].shape[2:] im_shape = data['im_shape'] scale_factor = data['scale_factor'] timer.tic() if not use_detector: dets = dets_list[frame_id] bbox_tlwh = paddle.to_tensor(dets['bbox'], dtype='float32') pred_scores = paddle.to_tensor(dets['score'], dtype='float32') if pred_scores < draw_threshold: continue if bbox_tlwh.shape[0] > 0: pred_bboxes = paddle.concat( (bbox_tlwh[:, 0:2], bbox_tlwh[:, 2:4] + bbox_tlwh[:, 0:2]), axis=1) else: pred_bboxes = [] pred_scores = [] else: outs = self.model.detector(data) if outs['bbox_num'] > 0: pred_bboxes = scale_coords(outs['bbox'][:, 2:], input_shape, im_shape, scale_factor) pred_scores = outs['bbox'][:, 1:2] else: pred_bboxes = [] pred_scores = [] pred_bboxes = clip_box(pred_bboxes, input_shape, im_shape, scale_factor) bbox_tlwh = paddle.concat( (pred_bboxes[:, 0:2], pred_bboxes[:, 2:4] - pred_bboxes[:, 0:2] + 1), axis=1) crops, pred_scores = get_crops( pred_bboxes, ori_image, pred_scores, w=64, h=192) crops = paddle.to_tensor(crops) pred_scores = paddle.to_tensor(pred_scores) data.update({'crops': crops}) features = self.model(data) features = features.numpy() detections = [ Detection(tlwh, score, feat) for tlwh, score, feat in zip(bbox_tlwh, pred_scores, features) ] self.model.tracker.predict() online_targets = self.model.tracker.update(detections) online_tlwhs = [] online_scores = [] online_ids = [] for track in online_targets: if not track.is_confirmed() or track.time_since_update > 1: continue online_tlwhs.append(track.to_tlwh()) online_scores.append(1.0) online_ids.append(track.track_id) timer.toc() # save results results.append( (frame_id + 1, online_tlwhs, online_scores, online_ids)) self.save_results(data, frame_id, online_ids, online_tlwhs, online_scores, timer.average_time, show_image, save_dir) frame_id += 1 return results, frame_id, timer.average_time, timer.calls def mot_evaluate(self, data_root, seqs, output_dir, data_type='mot', model_type='JDE', save_images=False, save_videos=False, show_image=False, det_results_dir=''): if not os.path.exists(output_dir): os.makedirs(output_dir) result_root = os.path.join(output_dir, 'mot_results') if not os.path.exists(result_root): os.makedirs(result_root) assert data_type in ['mot', 'kitti'], \ "data_type should be 'mot' or 'kitti'" assert model_type in ['JDE', 'DeepSORT', 'FairMOT'], \ "model_type should be 'JDE', 'DeepSORT' or 'FairMOT'" # run tracking n_frame = 0 timer_avgs, timer_calls = [], [] for seq in seqs: if not os.path.isdir(os.path.join(data_root, seq)): continue infer_dir = os.path.join(data_root, seq, 'img1') seqinfo = os.path.join(data_root, seq, 'seqinfo.ini') if not os.path.exists(seqinfo) or not os.path.exists( infer_dir) or not os.path.isdir(infer_dir): continue save_dir = os.path.join(output_dir, 'mot_outputs', seq) if save_images or save_videos else None logger.info('start seq: {}'.format(seq)) images = self.get_infer_images(infer_dir) self.dataset.set_images(images) dataloader = create('EvalMOTReader')(self.dataset, 0) result_filename = os.path.join(result_root, '{}.txt'.format(seq)) meta_info = open(seqinfo).read() frame_rate = int(meta_info[meta_info.find('frameRate') + 10: meta_info.find('\nseqLength')]) with paddle.no_grad(): if model_type in ['JDE', 'FairMOT']: results, nf, ta, tc = self._eval_seq_jde( dataloader, save_dir=save_dir, show_image=show_image, frame_rate=frame_rate) elif model_type in ['DeepSORT']: results, nf, ta, tc = self._eval_seq_sde( dataloader, save_dir=save_dir, show_image=show_image, frame_rate=frame_rate, det_file=os.path.join(det_results_dir, '{}.txt'.format(seq))) else: raise ValueError(model_type) self.write_mot_results(result_filename, results, data_type) n_frame += nf timer_avgs.append(ta) timer_calls.append(tc) if save_videos: output_video_path = os.path.join(save_dir, '..', '{}_vis.mp4'.format(seq)) cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg {}'.format( save_dir, output_video_path) os.system(cmd_str) logger.info('Save video in {}.'.format(output_video_path)) logger.info('Evaluate seq: {}'.format(seq)) # update metrics for metric in self._metrics: metric.update(data_root, seq, data_type, result_root, result_filename) timer_avgs = np.asarray(timer_avgs) timer_calls = np.asarray(timer_calls) all_time = np.dot(timer_avgs, timer_calls) avg_time = all_time / np.sum(timer_calls) logger.info('Time elapsed: {:.2f} seconds, FPS: {:.2f}'.format( all_time, 1.0 / avg_time)) # accumulate metric to log out for metric in self._metrics: metric.accumulate() metric.log() # reset metric states for metric may performed multiple times self._reset_metrics() def get_infer_images(self, infer_dir): assert infer_dir is None or os.path.isdir(infer_dir), \ "{} is not a directory".format(infer_dir) images = set() assert os.path.isdir(infer_dir), \ "infer_dir {} is not a directory".format(infer_dir) exts = ['jpg', 'jpeg', 'png', 'bmp'] exts += [ext.upper() for ext in exts] for ext in exts: images.update(glob.glob('{}/*.{}'.format(infer_dir, ext))) images = list(images) images.sort() assert len(images) > 0, "no image found in {}".format(infer_dir) logger.info("Found {} inference images in total.".format(len(images))) return images def mot_predict(self, video_file, frame_rate, image_dir, output_dir, data_type='mot', model_type='JDE', save_images=False, save_videos=True, show_image=False, det_results_dir='', draw_threshold=0.5): assert video_file is not None or image_dir is not None, \ "--video_file or --image_dir should be set." assert video_file is None or os.path.isfile(video_file), \ "{} is not a file".format(video_file) assert image_dir is None or os.path.isdir(image_dir), \ "{} is not a directory".format(image_dir) if not os.path.exists(output_dir): os.makedirs(output_dir) result_root = os.path.join(output_dir, 'mot_results') if not os.path.exists(result_root): os.makedirs(result_root) assert data_type in ['mot', 'kitti'], \ "data_type should be 'mot' or 'kitti'" assert model_type in ['JDE', 'DeepSORT', 'FairMOT'], \ "model_type should be 'JDE', 'DeepSORT' or 'FairMOT'" # run tracking if video_file: seq = video_file.split('/')[-1].split('.')[0] self.dataset.set_video(video_file, frame_rate) logger.info('Starting tracking video {}'.format(video_file)) elif image_dir: seq = image_dir.split('/')[-1].split('.')[0] images = [ '{}/{}'.format(image_dir, x) for x in os.listdir(image_dir) ] images.sort() self.dataset.set_images(images) logger.info('Starting tracking folder {}, found {} images'.format( image_dir, len(images))) else: raise ValueError('--video_file or --image_dir should be set.') save_dir = os.path.join(output_dir, 'mot_outputs', seq) if save_images or save_videos else None dataloader = create('TestMOTReader')(self.dataset, 0) result_filename = os.path.join(result_root, '{}.txt'.format(seq)) if frame_rate == -1: frame_rate = self.dataset.frame_rate with paddle.no_grad(): if model_type in ['JDE', 'FairMOT']: results, nf, ta, tc = self._eval_seq_jde( dataloader, save_dir=save_dir, show_image=show_image, frame_rate=frame_rate, draw_threshold=draw_threshold) elif model_type in ['DeepSORT']: results, nf, ta, tc = self._eval_seq_sde( dataloader, save_dir=save_dir, show_image=show_image, frame_rate=frame_rate, det_file=os.path.join(det_results_dir, '{}.txt'.format(seq)), draw_threshold=draw_threshold) else: raise ValueError(model_type) self.write_mot_results(result_filename, results, data_type) if save_videos: output_video_path = os.path.join(save_dir, '..', '{}_vis.mp4'.format(seq)) cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg {}'.format( save_dir, output_video_path) os.system(cmd_str) logger.info('Save video in {}'.format(output_video_path)) def write_mot_results(self, filename, results, data_type='mot'): if data_type in ['mot', 'mcmot', 'lab']: save_format = '{frame},{id},{x1},{y1},{w},{h},{score},-1,-1,-1\n' elif data_type == 'kitti': save_format = '{frame} {id} car 0 0 -10 {x1} {y1} {x2} {y2} -10 -10 -10 -1000 -1000 -1000 -10\n' else: raise ValueError(data_type) with open(filename, 'w') as f: for frame_id, tlwhs, tscores, track_ids in results: if data_type == 'kitti': frame_id -= 1 for tlwh, score, track_id in zip(tlwhs, tscores, track_ids): if track_id < 0: continue x1, y1, w, h = tlwh x2, y2 = x1 + w, y1 + h line = save_format.format( frame=frame_id, id=track_id, x1=x1, y1=y1, x2=x2, y2=y2, w=w, h=h, score=score) f.write(line) logger.info('MOT results save in {}'.format(filename)) def save_results(self, data, frame_id, online_ids, online_tlwhs, online_scores, average_time, show_image, save_dir): if show_image or save_dir is not None: assert 'ori_image' in data img0 = data['ori_image'].numpy()[0] online_im = mot_vis.plot_tracking( img0, online_tlwhs, online_ids, online_scores, frame_id=frame_id, fps=1. / average_time) if show_image: cv2.imshow('online_im', online_im) if save_dir is not None: cv2.imwrite( os.path.join(save_dir, '{:05d}.jpg'.format(frame_id)), online_im)