import os import os.path as osp import cv2 import logging import argparse import motmetrics as mm from tracker.mot_tracker_kalman import AETracker from utils import visualization as vis from utils.log import logger from utils.timer import Timer from utils.evaluation import Evaluator import utils.datasets as datasets import torch def mkdirs(path): if os.path.exists(path): return os.makedirs(path) def write_results(filename, results, data_type): if data_type == 'mot': save_format = '{frame},{id},{x1},{y1},{w},{h},1,-1,-1,-1\n' elif data_type == 'kitti': save_format = '{frame} {id} pedestrian 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, track_ids in results: if data_type == 'kitti': frame_id -= 1 for tlwh, track_id in zip(tlwhs, 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) f.write(line) logger.info('save results to {}'.format(filename)) def eval_seq(opt, dataloader, data_type, result_filename, save_dir=None, show_image=True, frame_rate=30): if save_dir is not None: mkdirs(save_dir) tracker = AETracker(opt, frame_rate=frame_rate) timer = Timer() results = [] frame_id = 0 for path, img, img0 in dataloader: if frame_id % 20 == 0: logger.info('Processing frame {} ({:.2f} fps)'.format(frame_id, 1./max(1e-5, timer.average_time))) # run tracking timer.tic() blob = torch.from_numpy(img).cuda().unsqueeze(0) online_targets = tracker.update(blob, img0) online_tlwhs = [] online_ids = [] for t in online_targets: tlwh = t.tlwh tid = t.track_id vertical = tlwh[2] / tlwh[3] > 1.6 if tlwh[2] * tlwh[3] > opt.min_box_area and not vertical: online_tlwhs.append(tlwh) online_ids.append(tid) timer.toc() # save results results.append((frame_id + 1, online_tlwhs, online_ids)) if show_image or save_dir is not None: online_im = vis.plot_tracking(img0, online_tlwhs, online_ids, frame_id=frame_id, fps=1. / timer.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) frame_id += 1 # save results write_results(result_filename, results, data_type) return frame_id def main(opt, data_root='/data/MOT16/train', det_root=None, seqs=('MOT16-05',), exp_name='demo', save_image=False, show_image=True): logger.setLevel(logging.INFO) result_root = os.path.join(data_root, '..', 'results', exp_name) mkdirs(result_root) data_type = 'mot' # run tracking timer = Timer() accs = [] n_frame = 0 timer.tic() for seq in seqs: output_dir = os.path.join(data_root, '..','outputs', exp_name, seq) if save_image else None logger.info('start seq: {}'.format(seq)) dataloader = datasets.LoadImages(osp.join(data_root, seq, 'img1'), opt.img_size) result_filename = os.path.join(result_root, '{}.txt'.format(seq)) meta_info = open(os.path.join(data_root, seq, 'seqinfo.ini')).read() frame_rate = int(meta_info[meta_info.find('frameRate')+10:meta_info.find('\nseqLength')]) n_frame += eval_seq(opt, dataloader, data_type, result_filename, save_dir=output_dir, show_image=show_image, frame_rate=frame_rate) # eval logger.info('Evaluate seq: {}'.format(seq)) evaluator = Evaluator(data_root, seq, data_type) accs.append(evaluator.eval_file(result_filename)) timer.toc() logger.info('Time elapsed: {}, FPS {}'.format(timer.average_time, n_frame / timer.average_time)) # get summary # metrics = ['mota', 'num_switches', 'idp', 'idr', 'idf1', 'precision', 'recall'] metrics = mm.metrics.motchallenge_metrics mh = mm.metrics.create() summary = Evaluator.get_summary(accs, seqs, metrics) strsummary = mm.io.render_summary( summary, formatters=mh.formatters, namemap=mm.io.motchallenge_metric_names ) print(strsummary) Evaluator.save_summary(summary, os.path.join(result_root, 'summary_{}.xlsx'.format(exp_name))) if __name__ == '__main__': parser = argparse.ArgumentParser(prog='test.py') parser.add_argument('--batch-size', type=int, default=8, help='size of each image batch') parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path') parser.add_argument('--weights', type=str, default='weights/latest.pt', help='path to weights file') parser.add_argument('--img-size', type=int, default=(864,480), help='size of each image dimension') parser.add_argument('--iou-thres', type=float, default=0.5, help='iou threshold required to qualify as detected') parser.add_argument('--conf-thres', type=float, default=0.5, help='object confidence threshold') parser.add_argument('--nms-thres', type=float, default=0.4, help='iou threshold for non-maximum suppression') parser.add_argument('--min-box-area', type=float, default=200, help='filter out tiny boxes') parser.add_argument('--pixel-mean', type=float, default=[0,0,0], nargs='+', help='pixel mean') parser.add_argument('--track-buffer', type=int, default=30, help='tracking buffer') parser.add_argument('--test-mot16', action='store_true', help='tracking buffer') parser.add_argument('--save-images', action='store_true', help='save tracking results') opt = parser.parse_args() print(opt, end='\n\n') if not opt.test_mot16: seqs_str = '''CVPR19-01 CVPR19-02 CVPR19-03 CVPR19-05''' data_root = '/home/wangzd/datasets/MOT/MOT19/train' else: seqs_str = '''MOT16-01 MOT16-03 MOT16-06 MOT16-07 MOT16-08 MOT16-12 MOT16-14''' #seqs_str = 'MOT16-14' data_root = '/home/wangzd/datasets/MOT/MOT16/test' seqs = [seq.strip() for seq in seqs_str.split()] main(opt, data_root=data_root, seqs=seqs, exp_name='darknet53.864x480', show_image=False, save_image=opt.save_images)