tracker.py 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538
  1. # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from __future__ import absolute_import
  15. from __future__ import division
  16. from __future__ import print_function
  17. import os
  18. import cv2
  19. import glob
  20. import re
  21. import paddle
  22. import numpy as np
  23. import os.path as osp
  24. from collections import defaultdict
  25. from paddlex.ppdet.core.workspace import create
  26. from paddlex.ppdet.utils.checkpoint import load_weight, load_pretrain_weight
  27. from paddlex.ppdet.modeling.mot.utils import Detection, get_crops, scale_coords, clip_box
  28. from paddlex.ppdet.modeling.mot.utils import MOTTimer, load_det_results, write_mot_results, save_vis_results
  29. from paddlex.ppdet.metrics import Metric, MOTMetric, KITTIMOTMetric
  30. from paddlex.ppdet.metrics import MCMOTMetric
  31. from .callbacks import Callback, ComposeCallback
  32. from paddlex.ppdet.utils.logger import setup_logger
  33. logger = setup_logger(__name__)
  34. __all__ = ['Tracker']
  35. class Tracker(object):
  36. def __init__(self, cfg, mode='eval'):
  37. self.cfg = cfg
  38. assert mode.lower() in ['test', 'eval'], \
  39. "mode should be 'test' or 'eval'"
  40. self.mode = mode.lower()
  41. self.optimizer = None
  42. # build MOT data loader
  43. self.dataset = cfg['{}MOTDataset'.format(self.mode.capitalize())]
  44. # build model
  45. self.model = create(cfg.architecture)
  46. self.status = {}
  47. self.start_epoch = 0
  48. # initial default callbacks
  49. self._init_callbacks()
  50. # initial default metrics
  51. self._init_metrics()
  52. self._reset_metrics()
  53. def _init_callbacks(self):
  54. self._callbacks = []
  55. self._compose_callback = None
  56. def _init_metrics(self):
  57. if self.mode in ['test']:
  58. self._metrics = []
  59. return
  60. if self.cfg.metric == 'MOT':
  61. self._metrics = [MOTMetric(), ]
  62. elif self.cfg.metric == 'MCMOT':
  63. self._metrics = [MCMOTMetric(self.cfg.num_classes), ]
  64. elif self.cfg.metric == 'KITTI':
  65. self._metrics = [KITTIMOTMetric(), ]
  66. else:
  67. logger.warning("Metric not support for metric type {}".format(
  68. self.cfg.metric))
  69. self._metrics = []
  70. def _reset_metrics(self):
  71. for metric in self._metrics:
  72. metric.reset()
  73. def register_callbacks(self, callbacks):
  74. callbacks = [h for h in list(callbacks) if h is not None]
  75. for c in callbacks:
  76. assert isinstance(c, Callback), \
  77. "metrics shoule be instances of subclass of Metric"
  78. self._callbacks.extend(callbacks)
  79. self._compose_callback = ComposeCallback(self._callbacks)
  80. def register_metrics(self, metrics):
  81. metrics = [m for m in list(metrics) if m is not None]
  82. for m in metrics:
  83. assert isinstance(m, Metric), \
  84. "metrics shoule be instances of subclass of Metric"
  85. self._metrics.extend(metrics)
  86. def load_weights_jde(self, weights):
  87. load_weight(self.model, weights, self.optimizer)
  88. def load_weights_sde(self, det_weights, reid_weights):
  89. if self.model.detector:
  90. load_weight(self.model.detector, det_weights)
  91. load_weight(self.model.reid, reid_weights)
  92. else:
  93. load_weight(self.model.reid, reid_weights, self.optimizer)
  94. def _eval_seq_jde(self,
  95. dataloader,
  96. save_dir=None,
  97. show_image=False,
  98. frame_rate=30,
  99. draw_threshold=0):
  100. if save_dir:
  101. if not os.path.exists(save_dir): os.makedirs(save_dir)
  102. tracker = self.model.tracker
  103. tracker.max_time_lost = int(frame_rate / 30.0 * tracker.track_buffer)
  104. timer = MOTTimer()
  105. frame_id = 0
  106. self.status['mode'] = 'track'
  107. self.model.eval()
  108. results = defaultdict(list) # support single class and multi classes
  109. for step_id, data in enumerate(dataloader):
  110. self.status['step_id'] = step_id
  111. if frame_id % 40 == 0:
  112. logger.info('Processing frame {} ({:.2f} fps)'.format(
  113. frame_id, 1. / max(1e-5, timer.average_time)))
  114. # forward
  115. timer.tic()
  116. pred_dets, pred_embs = self.model(data)
  117. pred_dets, pred_embs = pred_dets.numpy(), pred_embs.numpy()
  118. online_targets_dict = self.model.tracker.update(pred_dets,
  119. pred_embs)
  120. online_tlwhs = defaultdict(list)
  121. online_scores = defaultdict(list)
  122. online_ids = defaultdict(list)
  123. for cls_id in range(self.cfg.num_classes):
  124. online_targets = online_targets_dict[cls_id]
  125. for t in online_targets:
  126. tlwh = t.tlwh
  127. tid = t.track_id
  128. tscore = t.score
  129. if tlwh[2] * tlwh[3] <= tracker.min_box_area: continue
  130. if tracker.vertical_ratio > 0 and tlwh[2] / tlwh[
  131. 3] > tracker.vertical_ratio:
  132. continue
  133. online_tlwhs[cls_id].append(tlwh)
  134. online_ids[cls_id].append(tid)
  135. online_scores[cls_id].append(tscore)
  136. # save results
  137. results[cls_id].append(
  138. (frame_id + 1, online_tlwhs[cls_id], online_scores[cls_id],
  139. online_ids[cls_id]))
  140. timer.toc()
  141. save_vis_results(data, frame_id, online_ids, online_tlwhs,
  142. online_scores, timer.average_time, show_image,
  143. save_dir, self.cfg.num_classes)
  144. frame_id += 1
  145. return results, frame_id, timer.average_time, timer.calls
  146. def _eval_seq_sde(self,
  147. dataloader,
  148. save_dir=None,
  149. show_image=False,
  150. frame_rate=30,
  151. seq_name='',
  152. scaled=False,
  153. det_file='',
  154. draw_threshold=0):
  155. if save_dir:
  156. if not os.path.exists(save_dir): os.makedirs(save_dir)
  157. use_detector = False if not self.model.detector else True
  158. timer = MOTTimer()
  159. results = defaultdict(list)
  160. frame_id = 0
  161. self.status['mode'] = 'track'
  162. self.model.eval()
  163. self.model.reid.eval()
  164. if not use_detector:
  165. dets_list = load_det_results(det_file, len(dataloader))
  166. logger.info('Finish loading detection results file {}.'.format(
  167. det_file))
  168. for step_id, data in enumerate(dataloader):
  169. self.status['step_id'] = step_id
  170. if frame_id % 40 == 0:
  171. logger.info('Processing frame {} ({:.2f} fps)'.format(
  172. frame_id, 1. / max(1e-5, timer.average_time)))
  173. ori_image = data['ori_image'] # [bs, H, W, 3]
  174. ori_image_shape = data['ori_image'].shape[1:3]
  175. # ori_image_shape: [H, W]
  176. input_shape = data['image'].shape[2:]
  177. # input_shape: [h, w], before data transforms, set in model config
  178. im_shape = data['im_shape'][0].numpy()
  179. # im_shape: [new_h, new_w], after data transforms
  180. scale_factor = data['scale_factor'][0].numpy()
  181. empty_detections = False
  182. # when it has no detected bboxes, will not inference reid model
  183. # and if visualize, use original image instead
  184. # forward
  185. timer.tic()
  186. if not use_detector:
  187. dets = dets_list[frame_id]
  188. bbox_tlwh = np.array(dets['bbox'], dtype='float32')
  189. if bbox_tlwh.shape[0] > 0:
  190. # detector outputs: pred_cls_ids, pred_scores, pred_bboxes
  191. pred_cls_ids = np.array(dets['cls_id'], dtype='float32')
  192. pred_scores = np.array(dets['score'], dtype='float32')
  193. pred_bboxes = np.concatenate(
  194. (bbox_tlwh[:, 0:2],
  195. bbox_tlwh[:, 2:4] + bbox_tlwh[:, 0:2]),
  196. axis=1)
  197. else:
  198. logger.warning(
  199. 'Frame {} has not object, try to modify score threshold.'.
  200. format(frame_id))
  201. empty_detections = True
  202. else:
  203. outs = self.model.detector(data)
  204. outs['bbox'] = outs['bbox'].numpy()
  205. outs['bbox_num'] = outs['bbox_num'].numpy()
  206. if outs['bbox_num'] > 0 and empty_detections == False:
  207. # detector outputs: pred_cls_ids, pred_scores, pred_bboxes
  208. pred_cls_ids = outs['bbox'][:, 0:1]
  209. pred_scores = outs['bbox'][:, 1:2]
  210. if not scaled:
  211. # Note: scaled=False only in JDE YOLOv3 or other detectors
  212. # with LetterBoxResize and JDEBBoxPostProcess.
  213. #
  214. # 'scaled' means whether the coords after detector outputs
  215. # have been scaled back to the original image, set True
  216. # in general detector, set False in JDE YOLOv3.
  217. pred_bboxes = scale_coords(outs['bbox'][:, 2:],
  218. input_shape, im_shape,
  219. scale_factor)
  220. else:
  221. pred_bboxes = outs['bbox'][:, 2:]
  222. else:
  223. logger.warning(
  224. 'Frame {} has not detected object, try to modify score threshold.'.
  225. format(frame_id))
  226. empty_detections = True
  227. if not empty_detections:
  228. pred_xyxys, keep_idx = clip_box(pred_bboxes, ori_image_shape)
  229. if len(keep_idx[0]) == 0:
  230. logger.warning(
  231. 'Frame {} has not detected object left after clip_box.'.
  232. format(frame_id))
  233. empty_detections = True
  234. if empty_detections:
  235. timer.toc()
  236. # if visualize, use original image instead
  237. online_ids, online_tlwhs, online_scores = None, None, None
  238. save_vis_results(data, frame_id, online_ids, online_tlwhs,
  239. online_scores, timer.average_time, show_image,
  240. save_dir, self.cfg.num_classes)
  241. frame_id += 1
  242. # thus will not inference reid model
  243. continue
  244. pred_scores = pred_scores[keep_idx[0]]
  245. pred_cls_ids = pred_cls_ids[keep_idx[0]]
  246. pred_tlwhs = np.concatenate(
  247. (pred_xyxys[:, 0:2],
  248. pred_xyxys[:, 2:4] - pred_xyxys[:, 0:2] + 1),
  249. axis=1)
  250. pred_dets = np.concatenate(
  251. (pred_tlwhs, pred_scores, pred_cls_ids), axis=1)
  252. tracker = self.model.tracker
  253. crops = get_crops(
  254. pred_xyxys,
  255. ori_image,
  256. w=tracker.input_size[0],
  257. h=tracker.input_size[1])
  258. crops = paddle.to_tensor(crops)
  259. data.update({'crops': crops})
  260. pred_embs = self.model(data).numpy()
  261. tracker.predict()
  262. online_targets = tracker.update(pred_dets, pred_embs)
  263. online_tlwhs, online_scores, online_ids = [], [], []
  264. for t in online_targets:
  265. if not t.is_confirmed() or t.time_since_update > 1:
  266. continue
  267. tlwh = t.to_tlwh()
  268. tscore = t.score
  269. tid = t.track_id
  270. if tscore < draw_threshold: continue
  271. if tlwh[2] * tlwh[3] <= tracker.min_box_area: continue
  272. if tracker.vertical_ratio > 0 and tlwh[2] / tlwh[
  273. 3] > tracker.vertical_ratio:
  274. continue
  275. online_tlwhs.append(tlwh)
  276. online_scores.append(tscore)
  277. online_ids.append(tid)
  278. timer.toc()
  279. # save results
  280. results[0].append(
  281. (frame_id + 1, online_tlwhs, online_scores, online_ids))
  282. save_vis_results(data, frame_id, online_ids, online_tlwhs,
  283. online_scores, timer.average_time, show_image,
  284. save_dir, self.cfg.num_classes)
  285. frame_id += 1
  286. return results, frame_id, timer.average_time, timer.calls
  287. def mot_evaluate(self,
  288. data_root,
  289. seqs,
  290. output_dir,
  291. data_type='mot',
  292. model_type='JDE',
  293. save_images=False,
  294. save_videos=False,
  295. show_image=False,
  296. scaled=False,
  297. det_results_dir=''):
  298. if not os.path.exists(output_dir): os.makedirs(output_dir)
  299. result_root = os.path.join(output_dir, 'mot_results')
  300. if not os.path.exists(result_root): os.makedirs(result_root)
  301. assert data_type in ['mot', 'mcmot', 'kitti'], \
  302. "data_type should be 'mot', 'mcmot' or 'kitti'"
  303. assert model_type in ['JDE', 'DeepSORT', 'FairMOT'], \
  304. "model_type should be 'JDE', 'DeepSORT' or 'FairMOT'"
  305. # run tracking
  306. n_frame = 0
  307. timer_avgs, timer_calls = [], []
  308. for seq in seqs:
  309. infer_dir = os.path.join(data_root, seq)
  310. if not os.path.exists(infer_dir) or not os.path.isdir(infer_dir):
  311. logger.warning("Seq {} error, {} has no images.".format(
  312. seq, infer_dir))
  313. continue
  314. if os.path.exists(os.path.join(infer_dir, 'img1')):
  315. infer_dir = os.path.join(infer_dir, 'img1')
  316. frame_rate = 30
  317. seqinfo = os.path.join(data_root, seq, 'seqinfo.ini')
  318. if os.path.exists(seqinfo):
  319. meta_info = open(seqinfo).read()
  320. frame_rate = int(meta_info[meta_info.find('frameRate') + 10:
  321. meta_info.find('\nseqLength')])
  322. save_dir = os.path.join(
  323. output_dir, 'mot_outputs',
  324. seq) if save_images or save_videos else None
  325. logger.info('start seq: {}'.format(seq))
  326. self.dataset.set_images(self.get_infer_images(infer_dir))
  327. dataloader = create('EvalMOTReader')(self.dataset, 0)
  328. result_filename = os.path.join(result_root, '{}.txt'.format(seq))
  329. with paddle.no_grad():
  330. if model_type in ['JDE', 'FairMOT']:
  331. results, nf, ta, tc = self._eval_seq_jde(
  332. dataloader,
  333. save_dir=save_dir,
  334. show_image=show_image,
  335. frame_rate=frame_rate)
  336. elif model_type in ['DeepSORT']:
  337. results, nf, ta, tc = self._eval_seq_sde(
  338. dataloader,
  339. save_dir=save_dir,
  340. show_image=show_image,
  341. frame_rate=frame_rate,
  342. seq_name=seq,
  343. scaled=scaled,
  344. det_file=os.path.join(det_results_dir,
  345. '{}.txt'.format(seq)))
  346. else:
  347. raise ValueError(model_type)
  348. write_mot_results(result_filename, results, data_type,
  349. self.cfg.num_classes)
  350. n_frame += nf
  351. timer_avgs.append(ta)
  352. timer_calls.append(tc)
  353. if save_videos:
  354. output_video_path = os.path.join(save_dir, '..',
  355. '{}_vis.mp4'.format(seq))
  356. cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg {}'.format(
  357. save_dir, output_video_path)
  358. os.system(cmd_str)
  359. logger.info('Save video in {}.'.format(output_video_path))
  360. logger.info('Evaluate seq: {}'.format(seq))
  361. # update metrics
  362. for metric in self._metrics:
  363. metric.update(data_root, seq, data_type, result_root,
  364. result_filename)
  365. timer_avgs = np.asarray(timer_avgs)
  366. timer_calls = np.asarray(timer_calls)
  367. all_time = np.dot(timer_avgs, timer_calls)
  368. avg_time = all_time / np.sum(timer_calls)
  369. logger.info('Time elapsed: {:.2f} seconds, FPS: {:.2f}'.format(
  370. all_time, 1.0 / avg_time))
  371. # accumulate metric to log out
  372. for metric in self._metrics:
  373. metric.accumulate()
  374. metric.log()
  375. # reset metric states for metric may performed multiple times
  376. self._reset_metrics()
  377. def get_infer_images(self, infer_dir):
  378. assert infer_dir is None or os.path.isdir(infer_dir), \
  379. "{} is not a directory".format(infer_dir)
  380. images = set()
  381. assert os.path.isdir(infer_dir), \
  382. "infer_dir {} is not a directory".format(infer_dir)
  383. exts = ['jpg', 'jpeg', 'png', 'bmp']
  384. exts += [ext.upper() for ext in exts]
  385. for ext in exts:
  386. images.update(glob.glob('{}/*.{}'.format(infer_dir, ext)))
  387. images = list(images)
  388. images.sort()
  389. assert len(images) > 0, "no image found in {}".format(infer_dir)
  390. logger.info("Found {} inference images in total.".format(len(images)))
  391. return images
  392. def mot_predict_seq(self,
  393. video_file,
  394. frame_rate,
  395. image_dir,
  396. output_dir,
  397. data_type='mot',
  398. model_type='JDE',
  399. save_images=False,
  400. save_videos=True,
  401. show_image=False,
  402. scaled=False,
  403. det_results_dir='',
  404. draw_threshold=0.5):
  405. assert video_file is not None or image_dir is not None, \
  406. "--video_file or --image_dir should be set."
  407. assert video_file is None or os.path.isfile(video_file), \
  408. "{} is not a file".format(video_file)
  409. assert image_dir is None or os.path.isdir(image_dir), \
  410. "{} is not a directory".format(image_dir)
  411. if not os.path.exists(output_dir): os.makedirs(output_dir)
  412. result_root = os.path.join(output_dir, 'mot_results')
  413. if not os.path.exists(result_root): os.makedirs(result_root)
  414. assert data_type in ['mot', 'mcmot', 'kitti'], \
  415. "data_type should be 'mot', 'mcmot' or 'kitti'"
  416. assert model_type in ['JDE', 'DeepSORT', 'FairMOT'], \
  417. "model_type should be 'JDE', 'DeepSORT' or 'FairMOT'"
  418. # run tracking
  419. if video_file:
  420. seq = video_file.split('/')[-1].split('.')[0]
  421. self.dataset.set_video(video_file, frame_rate)
  422. logger.info('Starting tracking video {}'.format(video_file))
  423. elif image_dir:
  424. seq = image_dir.split('/')[-1].split('.')[0]
  425. if os.path.exists(os.path.join(image_dir, 'img1')):
  426. image_dir = os.path.join(image_dir, 'img1')
  427. images = [
  428. '{}/{}'.format(image_dir, x) for x in os.listdir(image_dir)
  429. ]
  430. images.sort()
  431. self.dataset.set_images(images)
  432. logger.info('Starting tracking folder {}, found {} images'.format(
  433. image_dir, len(images)))
  434. else:
  435. raise ValueError('--video_file or --image_dir should be set.')
  436. save_dir = os.path.join(output_dir, 'mot_outputs',
  437. seq) if save_images or save_videos else None
  438. dataloader = create('TestMOTReader')(self.dataset, 0)
  439. result_filename = os.path.join(result_root, '{}.txt'.format(seq))
  440. if frame_rate == -1:
  441. frame_rate = self.dataset.frame_rate
  442. with paddle.no_grad():
  443. if model_type in ['JDE', 'FairMOT']:
  444. results, nf, ta, tc = self._eval_seq_jde(
  445. dataloader,
  446. save_dir=save_dir,
  447. show_image=show_image,
  448. frame_rate=frame_rate,
  449. draw_threshold=draw_threshold)
  450. elif model_type in ['DeepSORT']:
  451. results, nf, ta, tc = self._eval_seq_sde(
  452. dataloader,
  453. save_dir=save_dir,
  454. show_image=show_image,
  455. frame_rate=frame_rate,
  456. seq_name=seq,
  457. scaled=scaled,
  458. det_file=os.path.join(det_results_dir,
  459. '{}.txt'.format(seq)),
  460. draw_threshold=draw_threshold)
  461. else:
  462. raise ValueError(model_type)
  463. if save_videos:
  464. output_video_path = os.path.join(save_dir, '..',
  465. '{}_vis.mp4'.format(seq))
  466. cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg {}'.format(
  467. save_dir, output_video_path)
  468. os.system(cmd_str)
  469. logger.info('Save video in {}'.format(output_video_path))
  470. write_mot_results(result_filename, results, data_type,
  471. self.cfg.num_classes)