mot.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360
  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 cv2
  16. import numpy as np
  17. from collections import OrderedDict
  18. try:
  19. from collections.abc import Sequence
  20. except Exception:
  21. from collections import Sequence
  22. from .dataset import DetDataset, _make_dataset, _is_valid_file
  23. from paddlex.ppdet.core.workspace import register, serializable
  24. from paddlex.ppdet.utils.logger import setup_logger
  25. logger = setup_logger(__name__)
  26. @register
  27. @serializable
  28. class MOTDataSet(DetDataset):
  29. """
  30. Load dataset with MOT format.
  31. Args:
  32. dataset_dir (str): root directory for dataset.
  33. image_lists (str|list): mot data image lists, muiti-source mot dataset.
  34. data_fields (list): key name of data dictionary, at least have 'image'.
  35. sample_num (int): number of samples to load, -1 means all.
  36. Notes:
  37. MOT datasets root directory following this:
  38. dataset/mot
  39. |——————image_lists
  40. | |——————caltech.train
  41. | |——————caltech.val
  42. | |——————mot16.train
  43. | |——————mot17.train
  44. | ......
  45. |——————Caltech
  46. |——————MOT17
  47. |——————......
  48. All the MOT datasets have the following structure:
  49. Caltech
  50. |——————images
  51. | └——————00001.jpg
  52. | |—————— ...
  53. | └——————0000N.jpg
  54. └——————labels_with_ids
  55. └——————00001.txt
  56. |—————— ...
  57. └——————0000N.txt
  58. or
  59. MOT17
  60. |——————images
  61. | └——————train
  62. | └——————test
  63. └——————labels_with_ids
  64. └——————train
  65. """
  66. def __init__(self,
  67. dataset_dir=None,
  68. image_lists=[],
  69. data_fields=['image'],
  70. sample_num=-1):
  71. super(MOTDataSet, self).__init__(
  72. dataset_dir=dataset_dir,
  73. data_fields=data_fields,
  74. sample_num=sample_num)
  75. self.dataset_dir = dataset_dir
  76. self.image_lists = image_lists
  77. if isinstance(self.image_lists, str):
  78. self.image_lists = [self.image_lists]
  79. self.roidbs = None
  80. self.cname2cid = None
  81. def get_anno(self):
  82. if self.image_lists == []:
  83. return
  84. # only used to get categories and metric
  85. return os.path.join(self.dataset_dir, 'image_lists',
  86. self.image_lists[0])
  87. def parse_dataset(self):
  88. self.img_files = OrderedDict()
  89. self.img_start_index = OrderedDict()
  90. self.label_files = OrderedDict()
  91. self.tid_num = OrderedDict()
  92. self.tid_start_index = OrderedDict()
  93. img_index = 0
  94. for data_name in self.image_lists:
  95. # check every data image list
  96. image_lists_dir = os.path.join(self.dataset_dir, 'image_lists')
  97. assert os.path.isdir(image_lists_dir), \
  98. "The {} is not a directory.".format(image_lists_dir)
  99. list_path = os.path.join(image_lists_dir, data_name)
  100. assert os.path.exists(list_path), \
  101. "The list path {} does not exist.".format(list_path)
  102. # record img_files, filter out empty ones
  103. with open(list_path, 'r') as file:
  104. self.img_files[data_name] = file.readlines()
  105. self.img_files[data_name] = [
  106. os.path.join(self.dataset_dir, x.strip())
  107. for x in self.img_files[data_name]
  108. ]
  109. self.img_files[data_name] = list(
  110. filter(lambda x: len(x) > 0, self.img_files[data_name]))
  111. self.img_start_index[data_name] = img_index
  112. img_index += len(self.img_files[data_name])
  113. # record label_files
  114. self.label_files[data_name] = [
  115. x.replace('images', 'labels_with_ids').replace(
  116. '.png', '.txt').replace('.jpg', '.txt')
  117. for x in self.img_files[data_name]
  118. ]
  119. for data_name, label_paths in self.label_files.items():
  120. max_index = -1
  121. for lp in label_paths:
  122. lb = np.loadtxt(lp)
  123. if len(lb) < 1:
  124. continue
  125. if len(lb.shape) < 2:
  126. img_max = lb[1]
  127. else:
  128. img_max = np.max(lb[:, 1])
  129. if img_max > max_index:
  130. max_index = img_max
  131. self.tid_num[data_name] = int(max_index + 1)
  132. last_index = 0
  133. for i, (k, v) in enumerate(self.tid_num.items()):
  134. self.tid_start_index[k] = last_index
  135. last_index += v
  136. self.total_identities = int(last_index + 1)
  137. self.num_imgs_each_data = [len(x) for x in self.img_files.values()]
  138. self.total_imgs = sum(self.num_imgs_each_data)
  139. logger.info('=' * 80)
  140. logger.info('MOT dataset summary: ')
  141. logger.info(self.tid_num)
  142. logger.info('total images: {}'.format(self.total_imgs))
  143. logger.info('image start index: {}'.format(self.img_start_index))
  144. logger.info('total identities: {}'.format(self.total_identities))
  145. logger.info('identity start index: {}'.format(self.tid_start_index))
  146. logger.info('=' * 80)
  147. records = []
  148. cname2cid = mot_label()
  149. for img_index in range(self.total_imgs):
  150. for i, (k, v) in enumerate(self.img_start_index.items()):
  151. if img_index >= v:
  152. data_name = list(self.label_files.keys())[i]
  153. start_index = v
  154. img_file = self.img_files[data_name][img_index - start_index]
  155. lbl_file = self.label_files[data_name][img_index - start_index]
  156. if not os.path.exists(img_file):
  157. logger.warn('Illegal image file: {}, and it will be ignored'.
  158. format(img_file))
  159. continue
  160. if not os.path.isfile(lbl_file):
  161. logger.warn('Illegal label file: {}, and it will be ignored'.
  162. format(lbl_file))
  163. continue
  164. labels = np.loadtxt(lbl_file, dtype=np.float32).reshape(-1, 6)
  165. # each row in labels (N, 6) is [gt_class, gt_identity, cx, cy, w, h]
  166. cx, cy = labels[:, 2], labels[:, 3]
  167. w, h = labels[:, 4], labels[:, 5]
  168. gt_bbox = np.stack((cx, cy, w, h)).T.astype('float32')
  169. gt_class = labels[:, 0:1].astype('int32')
  170. gt_score = np.ones((len(labels), 1)).astype('float32')
  171. gt_ide = labels[:, 1:2].astype('int32')
  172. for i, _ in enumerate(gt_ide):
  173. if gt_ide[i] > -1:
  174. gt_ide[i] += self.tid_start_index[data_name]
  175. mot_rec = {
  176. 'im_file': img_file,
  177. 'im_id': img_index,
  178. } if 'image' in self.data_fields else {}
  179. gt_rec = {
  180. 'gt_class': gt_class,
  181. 'gt_score': gt_score,
  182. 'gt_bbox': gt_bbox,
  183. 'gt_ide': gt_ide,
  184. }
  185. for k, v in gt_rec.items():
  186. if k in self.data_fields:
  187. mot_rec[k] = v
  188. records.append(mot_rec)
  189. if self.sample_num > 0 and img_index >= self.sample_num:
  190. break
  191. assert len(records) > 0, 'not found any mot record in %s' % (
  192. self.image_lists)
  193. self.roidbs, self.cname2cid = records, cname2cid
  194. def mot_label():
  195. labels_map = {'person': 0}
  196. return labels_map
  197. @register
  198. @serializable
  199. class MOTImageFolder(DetDataset):
  200. def __init__(self,
  201. task,
  202. dataset_dir=None,
  203. data_root=None,
  204. image_dir=None,
  205. sample_num=-1,
  206. keep_ori_im=False,
  207. **kwargs):
  208. super(MOTImageFolder, self).__init__(
  209. dataset_dir, image_dir, sample_num=sample_num)
  210. self.task = task
  211. self.data_root = data_root
  212. self.keep_ori_im = keep_ori_im
  213. self._imid2path = {}
  214. self.roidbs = None
  215. def check_or_download_dataset(self):
  216. return
  217. def parse_dataset(self, ):
  218. if not self.roidbs:
  219. self.roidbs = self._load_images()
  220. def _parse(self):
  221. image_dir = self.image_dir
  222. if not isinstance(image_dir, Sequence):
  223. image_dir = [image_dir]
  224. images = []
  225. for im_dir in image_dir:
  226. if os.path.isdir(im_dir):
  227. im_dir = os.path.join(self.dataset_dir, im_dir)
  228. images.extend(_make_dataset(im_dir))
  229. elif os.path.isfile(im_dir) and _is_valid_file(im_dir):
  230. images.append(im_dir)
  231. return images
  232. def _load_images(self):
  233. images = self._parse()
  234. ct = 0
  235. records = []
  236. for image in images:
  237. assert image != '' and os.path.isfile(image), \
  238. "Image {} not found".format(image)
  239. if self.sample_num > 0 and ct >= self.sample_num:
  240. break
  241. rec = {'im_id': np.array([ct]), 'im_file': image}
  242. if self.keep_ori_im:
  243. rec.update({'keep_ori_im': 1})
  244. self._imid2path[ct] = image
  245. ct += 1
  246. records.append(rec)
  247. assert len(records) > 0, "No image file found"
  248. return records
  249. def get_imid2path(self):
  250. return self._imid2path
  251. def set_images(self, images):
  252. self.image_dir = images
  253. self.roidbs = self._load_images()
  254. def _is_valid_video(f, extensions=('.mp4', '.avi', '.mov', '.rmvb', 'flv')):
  255. return f.lower().endswith(extensions)
  256. @register
  257. @serializable
  258. class MOTVideoDataset(DetDataset):
  259. """
  260. Load MOT dataset with MOT format from video for inference.
  261. Args:
  262. video_file (str): path of the video file
  263. dataset_dir (str): root directory for dataset.
  264. keep_ori_im (bool): whether to keep original image, default False.
  265. Set True when used during MOT model inference while saving
  266. images or video, or used in DeepSORT.
  267. """
  268. def __init__(self,
  269. video_file='',
  270. dataset_dir=None,
  271. keep_ori_im=False,
  272. **kwargs):
  273. super(MOTVideoDataset, self).__init__(dataset_dir=dataset_dir)
  274. self.video_file = video_file
  275. self.dataset_dir = dataset_dir
  276. self.keep_ori_im = keep_ori_im
  277. self.roidbs = None
  278. def parse_dataset(self, ):
  279. if not self.roidbs:
  280. self.roidbs = self._load_video_images()
  281. def _load_video_images(self):
  282. self.cap = cv2.VideoCapture(self.video_file)
  283. self.vn = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
  284. self.frame_rate = int(round(self.cap.get(cv2.CAP_PROP_FPS)))
  285. logger.info('Length of the video: {:d} frames'.format(self.vn))
  286. res = True
  287. ct = 0
  288. records = []
  289. while res:
  290. res, img = self.cap.read()
  291. image = np.ascontiguousarray(img, dtype=np.float32)
  292. image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
  293. im_shape = image.shape
  294. rec = {
  295. 'im_id': np.array([ct]),
  296. 'image': image,
  297. 'h': im_shape[0],
  298. 'w': im_shape[1],
  299. 'im_shape': np.array(
  300. im_shape[:2], dtype=np.float32),
  301. 'scale_factor': np.array(
  302. [1., 1.], dtype=np.float32),
  303. }
  304. if self.keep_ori_im:
  305. rec.update({'ori_image': image})
  306. ct += 1
  307. records.append(rec)
  308. records = records[:-1]
  309. assert len(records) > 0, "No image file found"
  310. return records
  311. def set_video(self, video_file):
  312. self.video_file = video_file
  313. assert os.path.isfile(self.video_file) and _is_valid_video(self.video_file), \
  314. "wrong or unsupported file format: {}".format(self.video_file)
  315. self.roidbs = self._load_video_images()