voc.py 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204
  1. # Copyright (c) 2019 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 numpy as np
  16. import xml.etree.ElementTree as ET
  17. from paddlex.ppdet.core.workspace import register, serializable
  18. from .dataset import DetDataset
  19. from paddlex.ppdet.utils.logger import setup_logger
  20. logger = setup_logger(__name__)
  21. @register
  22. @serializable
  23. class VOCDataSet(DetDataset):
  24. """
  25. Load dataset with PascalVOC format.
  26. Notes:
  27. `anno_path` must contains xml file and image file path for annotations.
  28. Args:
  29. dataset_dir (str): root directory for dataset.
  30. image_dir (str): directory for images.
  31. anno_path (str): voc annotation file path.
  32. data_fields (list): key name of data dictionary, at least have 'image'.
  33. sample_num (int): number of samples to load, -1 means all.
  34. label_list (str): if use_default_label is False, will load
  35. mapping between category and class index.
  36. """
  37. def __init__(self,
  38. dataset_dir=None,
  39. image_dir=None,
  40. anno_path=None,
  41. data_fields=['image'],
  42. sample_num=-1,
  43. label_list=None):
  44. super(VOCDataSet, self).__init__(
  45. dataset_dir=dataset_dir,
  46. image_dir=image_dir,
  47. anno_path=anno_path,
  48. data_fields=data_fields,
  49. sample_num=sample_num)
  50. self.label_list = label_list
  51. def parse_dataset(self, ):
  52. anno_path = os.path.join(self.dataset_dir, self.anno_path)
  53. image_dir = os.path.join(self.dataset_dir, self.image_dir)
  54. # mapping category name to class id
  55. # first_class:0, second_class:1, ...
  56. records = []
  57. ct = 0
  58. cname2cid = {}
  59. if self.label_list:
  60. label_path = os.path.join(self.dataset_dir, self.label_list)
  61. if not os.path.exists(label_path):
  62. raise ValueError("label_list {} does not exists".format(
  63. label_path))
  64. with open(label_path, 'r') as fr:
  65. label_id = 0
  66. for line in fr.readlines():
  67. cname2cid[line.strip()] = label_id
  68. label_id += 1
  69. else:
  70. cname2cid = pascalvoc_label()
  71. with open(anno_path, 'r') as fr:
  72. while True:
  73. line = fr.readline()
  74. if not line:
  75. break
  76. img_file, xml_file = [os.path.join(image_dir, x) \
  77. for x in line.strip().split()[:2]]
  78. if not os.path.exists(img_file):
  79. logger.warn(
  80. 'Illegal image file: {}, and it will be ignored'.
  81. format(img_file))
  82. continue
  83. if not os.path.isfile(xml_file):
  84. logger.warn('Illegal xml file: {}, and it will be ignored'.
  85. format(xml_file))
  86. continue
  87. tree = ET.parse(xml_file)
  88. if tree.find('id') is None:
  89. im_id = np.array([ct])
  90. else:
  91. im_id = np.array([int(tree.find('id').text)])
  92. objs = tree.findall('object')
  93. im_w = float(tree.find('size').find('width').text)
  94. im_h = float(tree.find('size').find('height').text)
  95. if im_w < 0 or im_h < 0:
  96. logger.warn(
  97. 'Illegal width: {} or height: {} in annotation, '
  98. 'and {} will be ignored'.format(im_w, im_h, xml_file))
  99. continue
  100. gt_bbox = []
  101. gt_class = []
  102. gt_score = []
  103. difficult = []
  104. for i, obj in enumerate(objs):
  105. cname = obj.find('name').text
  106. # user dataset may not contain difficult field
  107. _difficult = obj.find('difficult')
  108. _difficult = int(
  109. _difficult.text) if _difficult is not None else 0
  110. x1 = float(obj.find('bndbox').find('xmin').text)
  111. y1 = float(obj.find('bndbox').find('ymin').text)
  112. x2 = float(obj.find('bndbox').find('xmax').text)
  113. y2 = float(obj.find('bndbox').find('ymax').text)
  114. x1 = max(0, x1)
  115. y1 = max(0, y1)
  116. x2 = min(im_w - 1, x2)
  117. y2 = min(im_h - 1, y2)
  118. if x2 > x1 and y2 > y1:
  119. gt_bbox.append([x1, y1, x2, y2])
  120. gt_class.append([cname2cid[cname]])
  121. gt_score.append([1.])
  122. difficult.append([_difficult])
  123. else:
  124. logger.warn(
  125. 'Found an invalid bbox in annotations: xml_file: {}'
  126. ', x1: {}, y1: {}, x2: {}, y2: {}.'.format(
  127. xml_file, x1, y1, x2, y2))
  128. gt_bbox = np.array(gt_bbox).astype('float32')
  129. gt_class = np.array(gt_class).astype('int32')
  130. gt_score = np.array(gt_score).astype('float32')
  131. difficult = np.array(difficult).astype('int32')
  132. voc_rec = {
  133. 'im_file': img_file,
  134. 'im_id': im_id,
  135. 'h': im_h,
  136. 'w': im_w
  137. } if 'image' in self.data_fields else {}
  138. gt_rec = {
  139. 'gt_class': gt_class,
  140. 'gt_score': gt_score,
  141. 'gt_bbox': gt_bbox,
  142. 'difficult': difficult
  143. }
  144. for k, v in gt_rec.items():
  145. if k in self.data_fields:
  146. voc_rec[k] = v
  147. if len(objs) != 0:
  148. records.append(voc_rec)
  149. ct += 1
  150. if self.sample_num > 0 and ct >= self.sample_num:
  151. break
  152. assert len(records) > 0, 'not found any voc record in %s' % (
  153. self.anno_path)
  154. logger.debug('{} samples in file {}'.format(ct, anno_path))
  155. self.roidbs, self.cname2cid = records, cname2cid
  156. def get_label_list(self):
  157. return os.path.join(self.dataset_dir, self.label_list)
  158. def pascalvoc_label():
  159. labels_map = {
  160. 'aeroplane': 0,
  161. 'bicycle': 1,
  162. 'bird': 2,
  163. 'boat': 3,
  164. 'bottle': 4,
  165. 'bus': 5,
  166. 'car': 6,
  167. 'cat': 7,
  168. 'chair': 8,
  169. 'cow': 9,
  170. 'diningtable': 10,
  171. 'dog': 11,
  172. 'horse': 12,
  173. 'motorbike': 13,
  174. 'person': 14,
  175. 'pottedplant': 15,
  176. 'sheep': 16,
  177. 'sofa': 17,
  178. 'train': 18,
  179. 'tvmonitor': 19
  180. }
  181. return labels_map