171 lines
6.6 KiB
Python
171 lines
6.6 KiB
Python
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)
|
|
|