245 lines
9.9 KiB
Python
245 lines
9.9 KiB
Python
import os
|
|
import os.path as osp
|
|
import pickle
|
|
import cv2
|
|
import logging
|
|
import argparse
|
|
|
|
from tqdm.auto import tqdm
|
|
import motmetrics as mm
|
|
|
|
import torch
|
|
from tracker.multitracker import JDETracker
|
|
from utils import visualization as vis
|
|
from utils.log import logger
|
|
from utils.timer import Timer
|
|
from utils.evaluation import Evaluator
|
|
from utils.parse_config import parse_model_cfg
|
|
import utils.datasets as datasets
|
|
from utils.utils import *
|
|
|
|
|
|
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, save_img=False, save_figures=False, show_image=True, frame_rate=30):
|
|
'''
|
|
Processes the video sequence given and provides the output of tracking result (write the results in video file)
|
|
|
|
It uses JDE model for getting information about the online targets present.
|
|
|
|
Parameters
|
|
----------
|
|
opt : Namespace
|
|
Contains information passed as commandline arguments.
|
|
|
|
dataloader : LoadVideo
|
|
Instance of LoadVideo class used for fetching the image sequence and associated data.
|
|
|
|
data_type : String
|
|
Type of dataset corresponding(similar) to the given video.
|
|
|
|
result_filename : String
|
|
The name(path) of the file for storing results.
|
|
|
|
save_dir : String
|
|
Path to the folder for storing the frames containing bounding box information (Result frames). If given, featuers will be save there as pickle
|
|
save_figures : bool
|
|
If set, individual crops of all embedded figures will be saved
|
|
|
|
show_image : bool
|
|
Option for shhowing individial frames during run-time.
|
|
|
|
frame_rate : int
|
|
Frame-rate of the given video.
|
|
|
|
Returns
|
|
-------
|
|
(Returns are not significant here)
|
|
frame_id : int
|
|
Sequence number of the last sequence
|
|
'''
|
|
|
|
if save_dir:
|
|
mkdir_if_missing(save_dir)
|
|
tracker = JDETracker(opt, frame_rate=frame_rate)
|
|
timer = Timer()
|
|
results = []
|
|
frame_id = -1
|
|
for path, img, img0 in tqdm(dataloader):
|
|
frame_id += 1
|
|
# if frame_id % 20 == 0:
|
|
# logger.info('Processing frame {} ({:.2f} fps)'.format(frame_id, 1./max(1e-5, timer.average_time)))
|
|
frame_pickle_fn = os.path.join(save_dir, f'{frame_id:05d}.pcl')
|
|
if os.path.exists(frame_pickle_fn):
|
|
continue
|
|
# run tracking
|
|
timer.tic()
|
|
blob = torch.from_numpy(img).cuda().unsqueeze(0)
|
|
# online targets: all tartgets that are not timed out
|
|
# frame_embeddings: the embeddings of objects visible only in the current frame
|
|
online_targets, frame_embeddings = tracker.update(blob, img0)
|
|
online_tlwhs = []
|
|
online_ids = []
|
|
# for t in online_targets:
|
|
for t in tracker.tracked_stracks:
|
|
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:
|
|
base_fn = os.path.join(save_dir, '{:05d}'.format(frame_id))
|
|
if save_img:
|
|
cv2.imwrite(base_fn+'.jpg', online_im, [cv2.IMWRITE_JPEG_QUALITY, 100])
|
|
if save_figures:
|
|
for i, fe in enumerate(frame_embeddings):
|
|
tlwh, curr_feat = fe
|
|
x,y,w,h = round(tlwh[0]), round(tlwh[1]), round(tlwh[2]), round(tlwh[3])
|
|
# print(x,y,w,h, tlwh)
|
|
crop_img = img0[y:y+h, x:x+w]
|
|
cv2.imwrite(f'{base_fn}-{i}.jpg', crop_img)
|
|
with open(os.path.join(save_dir, f'{frame_id:05d}-{i}.pcl'), 'wb') as fp:
|
|
pickle.dump(fe, fp)
|
|
|
|
with open(frame_pickle_fn, 'wb') as fp:
|
|
pickle.dump(frame_embeddings, fp)
|
|
|
|
|
|
# save results
|
|
if result_filename is not None:
|
|
write_results(result_filename, results, data_type)
|
|
|
|
return frame_id, timer.average_time, timer.calls
|
|
|
|
|
|
def main(opt, data_root='/data/MOT16/train', det_root=None, seqs=('MOT16-05',), exp_name='demo',
|
|
save_images=False, save_videos=False, show_image=True):
|
|
logger.setLevel(logging.INFO)
|
|
result_root = os.path.join(data_root, '..', 'results', exp_name)
|
|
mkdir_if_missing(result_root)
|
|
data_type = 'mot'
|
|
|
|
# Read config
|
|
cfg_dict = parse_model_cfg(opt.cfg)
|
|
opt.img_size = [int(cfg_dict[0]['width']), int(cfg_dict[0]['height'])]
|
|
|
|
# run tracking
|
|
accs = []
|
|
n_frame = 0
|
|
timer_avgs, timer_calls = [], []
|
|
for seq in seqs:
|
|
output_dir = os.path.join(data_root, '..','outputs', exp_name, seq) if save_images or save_videos 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')])
|
|
nf, ta, tc = eval_seq(opt, dataloader, data_type, result_filename,
|
|
save_dir=output_dir, show_image=show_image, frame_rate=frame_rate)
|
|
n_frame += nf
|
|
timer_avgs.append(ta)
|
|
timer_calls.append(tc)
|
|
|
|
# eval
|
|
logger.info('Evaluate seq: {}'.format(seq))
|
|
evaluator = Evaluator(data_root, seq, data_type)
|
|
accs.append(evaluator.eval_file(result_filename))
|
|
if save_videos:
|
|
output_video_path = osp.join(output_dir, '{}.mp4'.format(seq))
|
|
cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg -c:v copy {}'.format(output_dir, output_video_path)
|
|
os.system(cmd_str)
|
|
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))
|
|
|
|
# get summary
|
|
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='track.py')
|
|
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('--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('--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 (image)')
|
|
parser.add_argument('--save-videos', action='store_true', help='save tracking results (video)')
|
|
opt = parser.parse_args()
|
|
print(opt, end='\n\n')
|
|
|
|
if not opt.test_mot16:
|
|
seqs_str = '''MOT17-02-SDP
|
|
MOT17-04-SDP
|
|
MOT17-05-SDP
|
|
MOT17-09-SDP
|
|
MOT17-10-SDP
|
|
MOT17-11-SDP
|
|
MOT17-13-SDP
|
|
'''
|
|
data_root = '/home/wangzd/datasets/MOT/MOT17/images/train'
|
|
else:
|
|
seqs_str = '''MOT16-01
|
|
MOT16-03
|
|
MOT16-06
|
|
MOT16-07
|
|
MOT16-08
|
|
MOT16-12
|
|
MOT16-14'''
|
|
data_root = '/home/wangzd/datasets/MOT/MOT16/images/test'
|
|
seqs = [seq.strip() for seq in seqs_str.split()]
|
|
|
|
main(opt,
|
|
data_root=data_root,
|
|
seqs=seqs,
|
|
exp_name=opt.weights.split('/')[-2],
|
|
show_image=False,
|
|
save_images=opt.save_images,
|
|
save_videos=opt.save_videos)
|
|
|