mot.py 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632
  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. import os
  15. import sys
  16. import cv2
  17. import glob
  18. import numpy as np
  19. from collections import OrderedDict, defaultdict
  20. try:
  21. from collections.abc import Sequence
  22. except Exception:
  23. from collections import Sequence
  24. from .dataset import DetDataset, _make_dataset, _is_valid_file
  25. from ppdet.core.workspace import register, serializable
  26. from ppdet.utils.logger import setup_logger
  27. logger = setup_logger(__name__)
  28. @register
  29. @serializable
  30. class MOTDataSet(DetDataset):
  31. """
  32. Load dataset with MOT format, only support single class MOT.
  33. Args:
  34. dataset_dir (str): root directory for dataset.
  35. image_lists (str|list): mot data image lists, muiti-source mot dataset.
  36. data_fields (list): key name of data dictionary, at least have 'image'.
  37. sample_num (int): number of samples to load, -1 means all.
  38. Notes:
  39. MOT datasets root directory following this:
  40. dataset/mot
  41. |——————image_lists
  42. | |——————caltech.train
  43. | |——————caltech.val
  44. | |——————mot16.train
  45. | |——————mot17.train
  46. | ......
  47. |——————Caltech
  48. |——————MOT17
  49. |——————......
  50. All the MOT datasets have the following structure:
  51. Caltech
  52. |——————images
  53. | └——————00001.jpg
  54. | |—————— ...
  55. | └——————0000N.jpg
  56. └——————labels_with_ids
  57. └——————00001.txt
  58. |—————— ...
  59. └——————0000N.txt
  60. or
  61. MOT17
  62. |——————images
  63. | └——————train
  64. | └——————test
  65. └——————labels_with_ids
  66. └——————train
  67. """
  68. def __init__(self,
  69. dataset_dir=None,
  70. image_lists=[],
  71. data_fields=['image'],
  72. sample_num=-1):
  73. super(MOTDataSet, self).__init__(
  74. dataset_dir=dataset_dir,
  75. data_fields=data_fields,
  76. sample_num=sample_num)
  77. self.dataset_dir = dataset_dir
  78. self.image_lists = image_lists
  79. if isinstance(self.image_lists, str):
  80. self.image_lists = [self.image_lists]
  81. self.roidbs = None
  82. self.cname2cid = None
  83. def get_anno(self):
  84. if self.image_lists == []:
  85. return
  86. # only used to get categories and metric
  87. # only check first data, but the label_list of all data should be same.
  88. first_mot_data = self.image_lists[0].split('.')[0]
  89. anno_file = os.path.join(self.dataset_dir, first_mot_data, 'label_list.txt')
  90. return anno_file
  91. def parse_dataset(self):
  92. self.img_files = OrderedDict()
  93. self.img_start_index = OrderedDict()
  94. self.label_files = OrderedDict()
  95. self.tid_num = OrderedDict()
  96. self.tid_start_index = OrderedDict()
  97. img_index = 0
  98. for data_name in self.image_lists:
  99. # check every data image list
  100. image_lists_dir = os.path.join(self.dataset_dir, 'image_lists')
  101. assert os.path.isdir(image_lists_dir), \
  102. "The {} is not a directory.".format(image_lists_dir)
  103. list_path = os.path.join(image_lists_dir, data_name)
  104. assert os.path.exists(list_path), \
  105. "The list path {} does not exist.".format(list_path)
  106. # record img_files, filter out empty ones
  107. with open(list_path, 'r') as file:
  108. self.img_files[data_name] = file.readlines()
  109. self.img_files[data_name] = [
  110. os.path.join(self.dataset_dir, x.strip())
  111. for x in self.img_files[data_name]
  112. ]
  113. self.img_files[data_name] = list(
  114. filter(lambda x: len(x) > 0, self.img_files[data_name]))
  115. self.img_start_index[data_name] = img_index
  116. img_index += len(self.img_files[data_name])
  117. # record label_files
  118. self.label_files[data_name] = [
  119. x.replace('images', 'labels_with_ids').replace(
  120. '.png', '.txt').replace('.jpg', '.txt')
  121. for x in self.img_files[data_name]
  122. ]
  123. for data_name, label_paths in self.label_files.items():
  124. max_index = -1
  125. for lp in label_paths:
  126. lb = np.loadtxt(lp)
  127. if len(lb) < 1:
  128. continue
  129. if len(lb.shape) < 2:
  130. img_max = lb[1]
  131. else:
  132. img_max = np.max(lb[:, 1])
  133. if img_max > max_index:
  134. max_index = img_max
  135. self.tid_num[data_name] = int(max_index + 1)
  136. last_index = 0
  137. for i, (k, v) in enumerate(self.tid_num.items()):
  138. self.tid_start_index[k] = last_index
  139. last_index += v
  140. self.num_identities_dict = defaultdict(int)
  141. self.num_identities_dict[0] = int(last_index + 1) # single class
  142. self.num_imgs_each_data = [len(x) for x in self.img_files.values()]
  143. self.total_imgs = sum(self.num_imgs_each_data)
  144. logger.info('MOT dataset summary: ')
  145. logger.info(self.tid_num)
  146. logger.info('Total images: {}'.format(self.total_imgs))
  147. logger.info('Image start index: {}'.format(self.img_start_index))
  148. logger.info('Total identities: {}'.format(self.num_identities_dict[0]))
  149. logger.info('Identity start index: {}'.format(self.tid_start_index))
  150. records = []
  151. cname2cid = mot_label()
  152. for img_index in range(self.total_imgs):
  153. for i, (k, v) in enumerate(self.img_start_index.items()):
  154. if img_index >= v:
  155. data_name = list(self.label_files.keys())[i]
  156. start_index = v
  157. img_file = self.img_files[data_name][img_index - start_index]
  158. lbl_file = self.label_files[data_name][img_index - start_index]
  159. if not os.path.exists(img_file):
  160. logger.warning('Illegal image file: {}, and it will be ignored'.
  161. format(img_file))
  162. continue
  163. if not os.path.isfile(lbl_file):
  164. logger.warning('Illegal label file: {}, and it will be ignored'.
  165. format(lbl_file))
  166. continue
  167. labels = np.loadtxt(lbl_file, dtype=np.float32).reshape(-1, 6)
  168. # each row in labels (N, 6) is [gt_class, gt_identity, cx, cy, w, h]
  169. cx, cy = labels[:, 2], labels[:, 3]
  170. w, h = labels[:, 4], labels[:, 5]
  171. gt_bbox = np.stack((cx, cy, w, h)).T.astype('float32')
  172. gt_class = labels[:, 0:1].astype('int32')
  173. gt_score = np.ones((len(labels), 1)).astype('float32')
  174. gt_ide = labels[:, 1:2].astype('int32')
  175. for i, _ in enumerate(gt_ide):
  176. if gt_ide[i] > -1:
  177. gt_ide[i] += self.tid_start_index[data_name]
  178. mot_rec = {
  179. 'im_file': img_file,
  180. 'im_id': img_index,
  181. } if 'image' in self.data_fields else {}
  182. gt_rec = {
  183. 'gt_class': gt_class,
  184. 'gt_score': gt_score,
  185. 'gt_bbox': gt_bbox,
  186. 'gt_ide': gt_ide,
  187. }
  188. for k, v in gt_rec.items():
  189. if k in self.data_fields:
  190. mot_rec[k] = v
  191. records.append(mot_rec)
  192. if self.sample_num > 0 and img_index >= self.sample_num:
  193. break
  194. assert len(records) > 0, 'not found any mot record in %s' % (
  195. self.image_lists)
  196. self.roidbs, self.cname2cid = records, cname2cid
  197. @register
  198. @serializable
  199. class MCMOTDataSet(DetDataset):
  200. """
  201. Load dataset with MOT format, support multi-class MOT.
  202. Args:
  203. dataset_dir (str): root directory for dataset.
  204. image_lists (list(str)): mcmot data image lists, muiti-source mcmot dataset.
  205. data_fields (list): key name of data dictionary, at least have 'image'.
  206. label_list (str): if use_default_label is False, will load
  207. mapping between category and class index.
  208. sample_num (int): number of samples to load, -1 means all.
  209. Notes:
  210. MCMOT datasets root directory following this:
  211. dataset/mot
  212. |——————image_lists
  213. | |——————visdrone_mcmot.train
  214. | |——————visdrone_mcmot.val
  215. visdrone_mcmot
  216. |——————images
  217. | └——————train
  218. | └——————val
  219. └——————labels_with_ids
  220. └——————train
  221. """
  222. def __init__(self,
  223. dataset_dir=None,
  224. image_lists=[],
  225. data_fields=['image'],
  226. label_list=None,
  227. sample_num=-1):
  228. super(MCMOTDataSet, self).__init__(
  229. dataset_dir=dataset_dir,
  230. data_fields=data_fields,
  231. sample_num=sample_num)
  232. self.dataset_dir = dataset_dir
  233. self.image_lists = image_lists
  234. if isinstance(self.image_lists, str):
  235. self.image_lists = [self.image_lists]
  236. self.label_list = label_list
  237. self.roidbs = None
  238. self.cname2cid = None
  239. def get_anno(self):
  240. if self.image_lists == []:
  241. return
  242. # only used to get categories and metric
  243. # only check first data, but the label_list of all data should be same.
  244. first_mot_data = self.image_lists[0].split('.')[0]
  245. anno_file = os.path.join(self.dataset_dir, first_mot_data, 'label_list.txt')
  246. return anno_file
  247. def parse_dataset(self):
  248. self.img_files = OrderedDict()
  249. self.img_start_index = OrderedDict()
  250. self.label_files = OrderedDict()
  251. self.tid_num = OrderedDict()
  252. self.tid_start_idx_of_cls_ids = defaultdict(dict) # for MCMOT
  253. img_index = 0
  254. for data_name in self.image_lists:
  255. # check every data image list
  256. image_lists_dir = os.path.join(self.dataset_dir, 'image_lists')
  257. assert os.path.isdir(image_lists_dir), \
  258. "The {} is not a directory.".format(image_lists_dir)
  259. list_path = os.path.join(image_lists_dir, data_name)
  260. assert os.path.exists(list_path), \
  261. "The list path {} does not exist.".format(list_path)
  262. # record img_files, filter out empty ones
  263. with open(list_path, 'r') as file:
  264. self.img_files[data_name] = file.readlines()
  265. self.img_files[data_name] = [
  266. os.path.join(self.dataset_dir, x.strip())
  267. for x in self.img_files[data_name]
  268. ]
  269. self.img_files[data_name] = list(
  270. filter(lambda x: len(x) > 0, self.img_files[data_name]))
  271. self.img_start_index[data_name] = img_index
  272. img_index += len(self.img_files[data_name])
  273. # record label_files
  274. self.label_files[data_name] = [
  275. x.replace('images', 'labels_with_ids').replace(
  276. '.png', '.txt').replace('.jpg', '.txt')
  277. for x in self.img_files[data_name]
  278. ]
  279. for data_name, label_paths in self.label_files.items():
  280. # using max_ids_dict rather than max_index
  281. max_ids_dict = defaultdict(int)
  282. for lp in label_paths:
  283. lb = np.loadtxt(lp)
  284. if len(lb) < 1:
  285. continue
  286. lb = lb.reshape(-1, 6)
  287. for item in lb:
  288. if item[1] > max_ids_dict[int(item[0])]:
  289. # item[0]: cls_id
  290. # item[1]: track id
  291. max_ids_dict[int(item[0])] = int(item[1])
  292. # track id number
  293. self.tid_num[data_name] = max_ids_dict
  294. last_idx_dict = defaultdict(int)
  295. for i, (k, v) in enumerate(self.tid_num.items()): # each sub dataset
  296. for cls_id, id_num in v.items(): # v is a max_ids_dict
  297. self.tid_start_idx_of_cls_ids[k][cls_id] = last_idx_dict[cls_id]
  298. last_idx_dict[cls_id] += id_num
  299. self.num_identities_dict = defaultdict(int)
  300. for k, v in last_idx_dict.items():
  301. self.num_identities_dict[k] = int(v) # total ids of each category
  302. self.num_imgs_each_data = [len(x) for x in self.img_files.values()]
  303. self.total_imgs = sum(self.num_imgs_each_data)
  304. # cname2cid and cid2cname
  305. cname2cid = {}
  306. if self.label_list is not None:
  307. # if use label_list for multi source mix dataset,
  308. # please make sure label_list in the first sub_dataset at least.
  309. sub_dataset = self.image_lists[0].split('.')[0]
  310. label_path = os.path.join(self.dataset_dir, sub_dataset,
  311. self.label_list)
  312. if not os.path.exists(label_path):
  313. logger.info(
  314. "Note: label_list {} does not exists, use VisDrone 10 classes labels as default.".
  315. format(label_path))
  316. cname2cid = visdrone_mcmot_label()
  317. else:
  318. with open(label_path, 'r') as fr:
  319. label_id = 0
  320. for line in fr.readlines():
  321. cname2cid[line.strip()] = label_id
  322. label_id += 1
  323. else:
  324. cname2cid = visdrone_mcmot_label()
  325. cid2cname = dict([(v, k) for (k, v) in cname2cid.items()])
  326. logger.info('MCMOT dataset summary: ')
  327. logger.info(self.tid_num)
  328. logger.info('Total images: {}'.format(self.total_imgs))
  329. logger.info('Image start index: {}'.format(self.img_start_index))
  330. logger.info('Total identities of each category: ')
  331. num_identities_dict = sorted(
  332. self.num_identities_dict.items(), key=lambda x: x[0])
  333. total_IDs_all_cats = 0
  334. for (k, v) in num_identities_dict:
  335. logger.info('Category {} [{}] has {} IDs.'.format(k, cid2cname[k],
  336. v))
  337. total_IDs_all_cats += v
  338. logger.info('Total identities of all categories: {}'.format(
  339. total_IDs_all_cats))
  340. logger.info('Identity start index of each category: ')
  341. for k, v in self.tid_start_idx_of_cls_ids.items():
  342. sorted_v = sorted(v.items(), key=lambda x: x[0])
  343. for (cls_id, start_idx) in sorted_v:
  344. logger.info('Start index of dataset {} category {:d} is {:d}'
  345. .format(k, cls_id, start_idx))
  346. records = []
  347. for img_index in range(self.total_imgs):
  348. for i, (k, v) in enumerate(self.img_start_index.items()):
  349. if img_index >= v:
  350. data_name = list(self.label_files.keys())[i]
  351. start_index = v
  352. img_file = self.img_files[data_name][img_index - start_index]
  353. lbl_file = self.label_files[data_name][img_index - start_index]
  354. if not os.path.exists(img_file):
  355. logger.warning('Illegal image file: {}, and it will be ignored'.
  356. format(img_file))
  357. continue
  358. if not os.path.isfile(lbl_file):
  359. logger.warning('Illegal label file: {}, and it will be ignored'.
  360. format(lbl_file))
  361. continue
  362. labels = np.loadtxt(lbl_file, dtype=np.float32).reshape(-1, 6)
  363. # each row in labels (N, 6) is [gt_class, gt_identity, cx, cy, w, h]
  364. cx, cy = labels[:, 2], labels[:, 3]
  365. w, h = labels[:, 4], labels[:, 5]
  366. gt_bbox = np.stack((cx, cy, w, h)).T.astype('float32')
  367. gt_class = labels[:, 0:1].astype('int32')
  368. gt_score = np.ones((len(labels), 1)).astype('float32')
  369. gt_ide = labels[:, 1:2].astype('int32')
  370. for i, _ in enumerate(gt_ide):
  371. if gt_ide[i] > -1:
  372. cls_id = int(gt_class[i])
  373. start_idx = self.tid_start_idx_of_cls_ids[data_name][cls_id]
  374. gt_ide[i] += start_idx
  375. mot_rec = {
  376. 'im_file': img_file,
  377. 'im_id': img_index,
  378. } if 'image' in self.data_fields else {}
  379. gt_rec = {
  380. 'gt_class': gt_class,
  381. 'gt_score': gt_score,
  382. 'gt_bbox': gt_bbox,
  383. 'gt_ide': gt_ide,
  384. }
  385. for k, v in gt_rec.items():
  386. if k in self.data_fields:
  387. mot_rec[k] = v
  388. records.append(mot_rec)
  389. if self.sample_num > 0 and img_index >= self.sample_num:
  390. break
  391. assert len(records) > 0, 'not found any mot record in %s' % (
  392. self.image_lists)
  393. self.roidbs, self.cname2cid = records, cname2cid
  394. @register
  395. @serializable
  396. class MOTImageFolder(DetDataset):
  397. """
  398. Load MOT dataset with MOT format from image folder or video .
  399. Args:
  400. video_file (str): path of the video file, default ''.
  401. frame_rate (int): frame rate of the video, use cv2 VideoCapture if not set.
  402. dataset_dir (str): root directory for dataset.
  403. keep_ori_im (bool): whether to keep original image, default False.
  404. Set True when used during MOT model inference while saving
  405. images or video, or used in DeepSORT.
  406. """
  407. def __init__(self,
  408. video_file=None,
  409. frame_rate=-1,
  410. dataset_dir=None,
  411. data_root=None,
  412. image_dir=None,
  413. sample_num=-1,
  414. keep_ori_im=False,
  415. anno_path=None,
  416. **kwargs):
  417. super(MOTImageFolder, self).__init__(
  418. dataset_dir, image_dir, sample_num=sample_num)
  419. self.video_file = video_file
  420. self.data_root = data_root
  421. self.keep_ori_im = keep_ori_im
  422. self._imid2path = {}
  423. self.roidbs = None
  424. self.frame_rate = frame_rate
  425. self.anno_path = anno_path
  426. def check_or_download_dataset(self):
  427. return
  428. def parse_dataset(self, ):
  429. if not self.roidbs:
  430. if self.video_file is None:
  431. self.frame_rate = 30 # set as default if infer image folder
  432. self.roidbs = self._load_images()
  433. else:
  434. self.roidbs = self._load_video_images()
  435. def _load_video_images(self):
  436. if self.frame_rate == -1:
  437. # if frame_rate is not set for video, use cv2.VideoCapture
  438. cap = cv2.VideoCapture(self.video_file)
  439. self.frame_rate = int(cap.get(cv2.CAP_PROP_FPS))
  440. extension = self.video_file.split('.')[-1]
  441. output_path = self.video_file.replace('.{}'.format(extension), '')
  442. frames_path = video2frames(self.video_file, output_path,
  443. self.frame_rate)
  444. self.video_frames = sorted(
  445. glob.glob(os.path.join(frames_path, '*.png')))
  446. self.video_length = len(self.video_frames)
  447. logger.info('Length of the video: {:d} frames.'.format(
  448. self.video_length))
  449. ct = 0
  450. records = []
  451. for image in self.video_frames:
  452. assert image != '' and os.path.isfile(image), \
  453. "Image {} not found".format(image)
  454. if self.sample_num > 0 and ct >= self.sample_num:
  455. break
  456. rec = {'im_id': np.array([ct]), 'im_file': image}
  457. if self.keep_ori_im:
  458. rec.update({'keep_ori_im': 1})
  459. self._imid2path[ct] = image
  460. ct += 1
  461. records.append(rec)
  462. assert len(records) > 0, "No image file found"
  463. return records
  464. def _find_images(self):
  465. image_dir = self.image_dir
  466. if not isinstance(image_dir, Sequence):
  467. image_dir = [image_dir]
  468. images = []
  469. for im_dir in image_dir:
  470. if os.path.isdir(im_dir):
  471. im_dir = os.path.join(self.dataset_dir, im_dir)
  472. images.extend(_make_dataset(im_dir))
  473. elif os.path.isfile(im_dir) and _is_valid_file(im_dir):
  474. images.append(im_dir)
  475. return images
  476. def _load_images(self):
  477. images = self._find_images()
  478. ct = 0
  479. records = []
  480. for image in images:
  481. assert image != '' and os.path.isfile(image), \
  482. "Image {} not found".format(image)
  483. if self.sample_num > 0 and ct >= self.sample_num:
  484. break
  485. rec = {'im_id': np.array([ct]), 'im_file': image}
  486. if self.keep_ori_im:
  487. rec.update({'keep_ori_im': 1})
  488. self._imid2path[ct] = image
  489. ct += 1
  490. records.append(rec)
  491. assert len(records) > 0, "No image file found"
  492. return records
  493. def get_imid2path(self):
  494. return self._imid2path
  495. def set_images(self, images):
  496. self.image_dir = images
  497. self.roidbs = self._load_images()
  498. def set_video(self, video_file, frame_rate):
  499. # update video_file and frame_rate by command line of tools/infer_mot.py
  500. self.video_file = video_file
  501. self.frame_rate = frame_rate
  502. assert os.path.isfile(self.video_file) and _is_valid_video(self.video_file), \
  503. "wrong or unsupported file format: {}".format(self.video_file)
  504. self.roidbs = self._load_video_images()
  505. def get_anno(self):
  506. return self.anno_path
  507. def _is_valid_video(f, extensions=('.mp4', '.avi', '.mov', '.rmvb', 'flv')):
  508. return f.lower().endswith(extensions)
  509. def video2frames(video_path, outpath, frame_rate, **kargs):
  510. def _dict2str(kargs):
  511. cmd_str = ''
  512. for k, v in kargs.items():
  513. cmd_str += (' ' + str(k) + ' ' + str(v))
  514. return cmd_str
  515. ffmpeg = ['ffmpeg ', ' -y -loglevel ', ' error ']
  516. vid_name = os.path.basename(video_path).split('.')[0]
  517. out_full_path = os.path.join(outpath, vid_name)
  518. if not os.path.exists(out_full_path):
  519. os.makedirs(out_full_path)
  520. # video file name
  521. outformat = os.path.join(out_full_path, '%08d.png')
  522. cmd = ffmpeg
  523. cmd = ffmpeg + [
  524. ' -i ', video_path, ' -r ', str(frame_rate), ' -f image2 ', outformat
  525. ]
  526. cmd = ''.join(cmd) + _dict2str(kargs)
  527. if os.system(cmd) != 0:
  528. raise RuntimeError('ffmpeg process video: {} error'.format(video_path))
  529. sys.exit(-1)
  530. sys.stdout.flush()
  531. return out_full_path
  532. def mot_label():
  533. labels_map = {'person': 0}
  534. return labels_map
  535. def visdrone_mcmot_label():
  536. labels_map = {
  537. 'pedestrian': 0,
  538. 'people': 1,
  539. 'bicycle': 2,
  540. 'car': 3,
  541. 'van': 4,
  542. 'truck': 5,
  543. 'tricycle': 6,
  544. 'awning-tricycle': 7,
  545. 'bus': 8,
  546. 'motor': 9,
  547. }
  548. return labels_map