voc.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341
  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. import copy
  16. import os.path as osp
  17. import random
  18. import re
  19. import numpy as np
  20. from collections import OrderedDict
  21. import xml.etree.ElementTree as ET
  22. from paddle.io import Dataset
  23. from paddlex.utils import logging, get_num_workers, get_encoding, path_normalization, is_pic
  24. from paddlex.cv.transforms import Decode, MixupImage
  25. from paddlex.tools import YOLOAnchorCluster
  26. class VOCDetection(Dataset):
  27. """读取PascalVOC格式的检测数据集,并对样本进行相应的处理。
  28. Args:
  29. data_dir (str): 数据集所在的目录路径。
  30. file_list (str): 描述数据集图片文件和对应标注文件的文件路径(文本内每行路径为相对data_dir的相对路)。
  31. label_list (str): 描述数据集包含的类别信息文件路径。
  32. transforms (paddlex.det.transforms): 数据集中每个样本的预处理/增强算子。
  33. num_workers (int|str): 数据集中样本在预处理过程中的线程或进程数。默认为'auto'。当设为'auto'时,根据
  34. 系统的实际CPU核数设置`num_workers`: 如果CPU核数的一半大于8,则`num_workers`为8,否则为CPU核数的
  35. 一半。
  36. shuffle (bool): 是否需要对数据集中样本打乱顺序。默认为False。
  37. """
  38. def __init__(self,
  39. data_dir,
  40. file_list,
  41. label_list,
  42. transforms=None,
  43. num_workers='auto',
  44. shuffle=False):
  45. # matplotlib.use() must be called *before* pylab, matplotlib.pyplot,
  46. # or matplotlib.backends is imported for the first time
  47. # pycocotools import matplotlib
  48. import matplotlib
  49. matplotlib.use('Agg')
  50. from pycocotools.coco import COCO
  51. super(VOCDetection, self).__init__()
  52. self.data_dir = data_dir
  53. self.data_fields = None
  54. self.transforms = copy.deepcopy(transforms)
  55. self.num_max_boxes = 50
  56. self.use_mix = False
  57. if self.transforms is not None:
  58. for op in self.transforms.transforms:
  59. if isinstance(op, MixupImage):
  60. self.mixup_op = copy.deepcopy(op)
  61. self.use_mix = True
  62. self.num_max_boxes *= 2
  63. break
  64. self.batch_transforms = None
  65. self.num_workers = get_num_workers(num_workers)
  66. self.shuffle = shuffle
  67. self.file_list = list()
  68. self.labels = list()
  69. annotations = dict()
  70. annotations['images'] = list()
  71. annotations['categories'] = list()
  72. annotations['annotations'] = list()
  73. cname2cid = OrderedDict()
  74. label_id = 0
  75. with open(label_list, 'r', encoding=get_encoding(label_list)) as f:
  76. for line in f.readlines():
  77. cname2cid[line.strip()] = label_id
  78. label_id += 1
  79. self.labels.append(line.strip())
  80. logging.info("Starting to read file list from dataset...")
  81. for k, v in cname2cid.items():
  82. annotations['categories'].append({
  83. 'supercategory': 'component',
  84. 'id': v + 1,
  85. 'name': k
  86. })
  87. ct = 0
  88. ann_ct = 0
  89. with open(file_list, 'r', encoding=get_encoding(file_list)) as f:
  90. while True:
  91. line = f.readline()
  92. if not line:
  93. break
  94. if len(line.strip().split()) > 2:
  95. raise Exception("A space is defined as the separator, "
  96. "but it exists in image or label name {}."
  97. .format(line))
  98. img_file, xml_file = [
  99. osp.join(data_dir, x) for x in line.strip().split()[:2]
  100. ]
  101. img_file = path_normalization(img_file)
  102. xml_file = path_normalization(xml_file)
  103. if not is_pic(img_file):
  104. continue
  105. if not osp.isfile(xml_file):
  106. continue
  107. if not osp.exists(img_file):
  108. logging.warning('The image file {} does not exist!'.format(
  109. img_file))
  110. continue
  111. if not osp.exists(xml_file):
  112. logging.warning('The annotation file {} does not exist!'.
  113. format(xml_file))
  114. continue
  115. tree = ET.parse(xml_file)
  116. if tree.find('id') is None:
  117. im_id = np.array([ct])
  118. else:
  119. ct = int(tree.find('id').text)
  120. im_id = np.array([int(tree.find('id').text)])
  121. pattern = re.compile('<object>', re.IGNORECASE)
  122. obj_match = pattern.findall(
  123. str(ET.tostringlist(tree.getroot())))
  124. if len(obj_match) == 0:
  125. continue
  126. obj_tag = obj_match[0][1:-1]
  127. objs = tree.findall(obj_tag)
  128. pattern = re.compile('<size>', re.IGNORECASE)
  129. size_tag = pattern.findall(
  130. str(ET.tostringlist(tree.getroot())))
  131. if len(size_tag) > 0:
  132. size_tag = size_tag[0][1:-1]
  133. size_element = tree.find(size_tag)
  134. pattern = re.compile('<width>', re.IGNORECASE)
  135. width_tag = pattern.findall(
  136. str(ET.tostringlist(size_element)))[0][1:-1]
  137. im_w = float(size_element.find(width_tag).text)
  138. pattern = re.compile('<height>', re.IGNORECASE)
  139. height_tag = pattern.findall(
  140. str(ET.tostringlist(size_element)))[0][1:-1]
  141. im_h = float(size_element.find(height_tag).text)
  142. else:
  143. im_w = 0
  144. im_h = 0
  145. gt_bbox = np.zeros((len(objs), 4), dtype=np.float32)
  146. gt_class = np.zeros((len(objs), 1), dtype=np.int32)
  147. gt_score = np.ones((len(objs), 1), dtype=np.float32)
  148. is_crowd = np.zeros((len(objs), 1), dtype=np.int32)
  149. difficult = np.zeros((len(objs), 1), dtype=np.int32)
  150. skipped_indices = list()
  151. for i, obj in enumerate(objs):
  152. pattern = re.compile('<name>', re.IGNORECASE)
  153. name_tag = pattern.findall(str(ET.tostringlist(obj)))[0][
  154. 1:-1]
  155. cname = obj.find(name_tag).text.strip()
  156. gt_class[i][0] = cname2cid[cname]
  157. pattern = re.compile('<difficult>', re.IGNORECASE)
  158. diff_tag = pattern.findall(str(ET.tostringlist(obj)))
  159. if len(diff_tag) == 0:
  160. _difficult = 0
  161. else:
  162. diff_tag = diff_tag[0][1:-1]
  163. try:
  164. _difficult = int(obj.find(diff_tag).text)
  165. except Exception:
  166. _difficult = 0
  167. pattern = re.compile('<bndbox>', re.IGNORECASE)
  168. box_tag = pattern.findall(str(ET.tostringlist(obj)))
  169. if len(box_tag) == 0:
  170. logging.warning(
  171. "There's no field '<bndbox>' in one of object, "
  172. "so this object will be ignored. xml file: {}".
  173. format(xml_file))
  174. continue
  175. box_tag = box_tag[0][1:-1]
  176. box_element = obj.find(box_tag)
  177. pattern = re.compile('<xmin>', re.IGNORECASE)
  178. xmin_tag = pattern.findall(
  179. str(ET.tostringlist(box_element)))[0][1:-1]
  180. x1 = float(box_element.find(xmin_tag).text)
  181. pattern = re.compile('<ymin>', re.IGNORECASE)
  182. ymin_tag = pattern.findall(
  183. str(ET.tostringlist(box_element)))[0][1:-1]
  184. y1 = float(box_element.find(ymin_tag).text)
  185. pattern = re.compile('<xmax>', re.IGNORECASE)
  186. xmax_tag = pattern.findall(
  187. str(ET.tostringlist(box_element)))[0][1:-1]
  188. x2 = float(box_element.find(xmax_tag).text)
  189. pattern = re.compile('<ymax>', re.IGNORECASE)
  190. ymax_tag = pattern.findall(
  191. str(ET.tostringlist(box_element)))[0][1:-1]
  192. y2 = float(box_element.find(ymax_tag).text)
  193. x1 = max(0, x1)
  194. y1 = max(0, y1)
  195. if im_w > 0.5 and im_h > 0.5:
  196. x2 = min(im_w - 1, x2)
  197. y2 = min(im_h - 1, y2)
  198. if not (x2 >= x1 and y2 >= y1):
  199. skipped_indices.append(i)
  200. logging.warning(
  201. "Bounding box for object {} does not satisfy x1 <= x2 and y1 <= y2, "
  202. "so this object is skipped".format(i))
  203. continue
  204. gt_bbox[i] = [x1, y1, x2, y2]
  205. is_crowd[i][0] = 0
  206. difficult[i][0] = _difficult
  207. annotations['annotations'].append({
  208. 'iscrowd': 0,
  209. 'image_id': int(im_id[0]),
  210. 'bbox': [x1, y1, x2 - x1, y2 - y1],
  211. 'area': float((x2 - x1) * (y2 - y1)),
  212. 'category_id': cname2cid[cname] + 1,
  213. 'id': ann_ct,
  214. 'difficult': _difficult
  215. })
  216. ann_ct += 1
  217. if skipped_indices:
  218. gt_bbox = np.delete(gt_bbox, skipped_indices, axis=0)
  219. gt_class = np.delete(gt_class, skipped_indices, axis=0)
  220. gt_score = np.delete(gt_score, skipped_indices, axis=0)
  221. is_crowd = np.delete(is_crowd, skipped_indices, axis=0)
  222. difficult = np.delete(difficult, skipped_indices, axis=0)
  223. im_info = {
  224. 'im_id': im_id,
  225. 'image_shape': np.array([im_h, im_w]).astype('int32'),
  226. }
  227. label_info = {
  228. 'is_crowd': is_crowd,
  229. 'gt_class': gt_class,
  230. 'gt_bbox': gt_bbox,
  231. 'gt_score': gt_score,
  232. 'difficult': difficult
  233. }
  234. if gt_bbox.size != 0:
  235. self.file_list.append({
  236. 'image': img_file,
  237. **
  238. im_info,
  239. **
  240. label_info
  241. })
  242. ct += 1
  243. annotations['images'].append({
  244. 'height': im_h,
  245. 'width': im_w,
  246. 'id': int(im_id[0]),
  247. 'file_name': osp.split(img_file)[1]
  248. })
  249. if self.use_mix:
  250. self.num_max_boxes = max(self.num_max_boxes, 2 * len(objs))
  251. else:
  252. self.num_max_boxes = max(self.num_max_boxes, len(objs))
  253. if not len(self.file_list) > 0:
  254. raise Exception('not found any voc record in %s' % (file_list))
  255. logging.info("{} samples in file {}".format(
  256. len(self.file_list), file_list))
  257. self.num_samples = len(self.file_list)
  258. self.coco_gt = COCO()
  259. self.coco_gt.dataset = annotations
  260. self.coco_gt.createIndex()
  261. self._epoch = 0
  262. def __getitem__(self, idx):
  263. sample = copy.deepcopy(self.file_list[idx])
  264. if self.data_fields is not None:
  265. sample = {k: sample[k] for k in self.data_fields}
  266. if self.use_mix and (self.mixup_op.mixup_epoch == -1 or
  267. self._epoch < self.mixup_op.mixup_epoch):
  268. if self.num_samples > 1:
  269. mix_idx = random.randint(1, self.num_samples - 1)
  270. mix_pos = (mix_idx + idx) % self.num_samples
  271. else:
  272. mix_pos = 0
  273. sample_mix = copy.deepcopy(self.file_list[mix_pos])
  274. if self.data_fields is not None:
  275. sample_mix = {k: sample_mix[k] for k in self.data_fields}
  276. sample = self.mixup_op(sample=[
  277. Decode(to_rgb=False)(sample), Decode(to_rgb=False)(sample_mix)
  278. ])
  279. sample = self.transforms(sample)
  280. return sample
  281. def __len__(self):
  282. return self.num_samples
  283. def set_epoch(self, epoch_id):
  284. self._epoch = epoch_id
  285. def cluster_yolo_anchor(self,
  286. num_anchors,
  287. image_size,
  288. cache=True,
  289. cache_path=None,
  290. iters=300,
  291. gen_iters=1000,
  292. thresh=.25):
  293. """
  294. Cluster YOLO anchors.
  295. Reference:
  296. https://github.com/ultralytics/yolov5/blob/master/utils/autoanchor.py
  297. Args:
  298. num_anchors (int): number of clusters
  299. image_size (list or int): [h, w], being an int means image height and image width are the same.
  300. cache (bool): whether using cache
  301. cache_path (str or None, optional): cache directory path. If None, use `data_dir` of dataset.
  302. iters (int, optional): iters of kmeans algorithm
  303. gen_iters (int, optional): iters of genetic algorithm
  304. threshold (float, optional): anchor scale threshold
  305. verbose (bool, optional): whether print results
  306. """
  307. if cache_path is None:
  308. cache_path = self.data_dir
  309. cluster = YOLOAnchorCluster(
  310. num_anchors=num_anchors,
  311. dataset=self,
  312. image_size=image_size,
  313. cache=cache,
  314. cache_path=cache_path,
  315. iters=iters,
  316. gen_iters=gen_iters,
  317. thresh=thresh)
  318. anchors = cluster()
  319. return anchors