mot.py 24 KB

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