coco.py 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162
  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 six
  18. import sys
  19. import numpy as np
  20. from paddlex.utils import logging, is_pic, get_num_workers
  21. from .voc import VOCDetection
  22. from paddlex.cv.transforms import MixupImage
  23. class CocoDetection(VOCDetection):
  24. """读取MSCOCO格式的检测数据集,并对样本进行相应的处理,该格式的数据集同样可以应用到实例分割模型的训练中。
  25. Args:
  26. data_dir (str): 数据集所在的目录路径。
  27. ann_file (str): 数据集的标注文件,为一个独立的json格式文件。
  28. transforms (paddlex.det.transforms): 数据集中每个样本的预处理/增强算子。
  29. num_workers (int|str): 数据集中样本在预处理过程中的线程或进程数。默认为'auto'。当设为'auto'时,根据
  30. 系统的实际CPU核数设置`num_workers`: 如果CPU核数的一半大于8,则`num_workers`为8,否则为CPU核数的一半。
  31. shuffle (bool): 是否需要对数据集中样本打乱顺序。默认为False。
  32. """
  33. def __init__(self,
  34. data_dir,
  35. ann_file,
  36. transforms=None,
  37. num_workers='auto',
  38. shuffle=False):
  39. # matplotlib.use() must be called *before* pylab, matplotlib.pyplot,
  40. # or matplotlib.backends is imported for the first time
  41. # pycocotools import matplotlib
  42. import matplotlib
  43. matplotlib.use('Agg')
  44. from pycocotools.coco import COCO
  45. try:
  46. import shapely.ops
  47. from shapely.geometry import Polygon, MultiPolygon, GeometryCollection
  48. except:
  49. six.reraise(*sys.exc_info())
  50. super(VOCDetection, self).__init__()
  51. self.data_fields = None
  52. self.transforms = copy.deepcopy(transforms)
  53. self.use_mix = False
  54. if self.transforms is not None:
  55. for op in self.transforms.transforms:
  56. if isinstance(op, MixupImage):
  57. self.mixup_op = copy.deepcopy(op)
  58. self.use_mix = True
  59. break
  60. self.batch_transforms = None
  61. self.num_workers = get_num_workers(num_workers)
  62. self.shuffle = shuffle
  63. self.file_list = list()
  64. self.labels = list()
  65. coco = COCO(ann_file)
  66. self.coco_gt = coco
  67. img_ids = coco.getImgIds()
  68. cat_ids = coco.getCatIds()
  69. catid2clsid = dict({catid: i for i, catid in enumerate(cat_ids)})
  70. cname2clsid = dict({
  71. coco.loadCats(catid)[0]['name']: clsid
  72. for catid, clsid in catid2clsid.items()
  73. })
  74. for label, cid in sorted(cname2clsid.items(), key=lambda d: d[1]):
  75. self.labels.append(label)
  76. logging.info("Starting to read file list from dataset...")
  77. for img_id in img_ids:
  78. img_anno = coco.loadImgs(img_id)[0]
  79. im_fname = osp.join(data_dir, img_anno['file_name'])
  80. if not is_pic(im_fname):
  81. continue
  82. im_w = float(img_anno['width'])
  83. im_h = float(img_anno['height'])
  84. ins_anno_ids = coco.getAnnIds(imgIds=img_id, iscrowd=False)
  85. instances = coco.loadAnns(ins_anno_ids)
  86. bboxes = []
  87. for inst in instances:
  88. x, y, box_w, box_h = inst['bbox']
  89. x1 = max(0, x)
  90. y1 = max(0, y)
  91. x2 = min(im_w - 1, x1 + max(0, box_w))
  92. y2 = min(im_h - 1, y1 + max(0, box_h))
  93. if inst['area'] > 0 and x2 >= x1 and y2 >= y1:
  94. inst['clean_bbox'] = [x1, y1, x2, y2]
  95. bboxes.append(inst)
  96. else:
  97. logging.warning(
  98. "Found an invalid bbox in annotations: "
  99. "im_id: {}, area: {} x1: {}, y1: {}, x2: {}, y2: {}."
  100. .format(img_id, float(inst['area']), x1, y1, x2, y2))
  101. num_bbox = len(bboxes)
  102. gt_bbox = np.zeros((num_bbox, 4), dtype=np.float32)
  103. gt_class = np.zeros((num_bbox, 1), dtype=np.int32)
  104. gt_score = np.ones((num_bbox, 1), dtype=np.float32)
  105. is_crowd = np.zeros((num_bbox, 1), dtype=np.int32)
  106. difficult = np.zeros((num_bbox, 1), dtype=np.int32)
  107. gt_poly = [None] * num_bbox
  108. has_segmentation = False
  109. for i, box in enumerate(bboxes):
  110. catid = box['category_id']
  111. gt_class[i][0] = catid2clsid[catid]
  112. gt_bbox[i, :] = box['clean_bbox']
  113. is_crowd[i][0] = box['iscrowd']
  114. if 'segmentation' in box and box['iscrowd'] == 1:
  115. gt_poly[i] = [[0.0, 0.0], ]
  116. elif 'segmentation' in box and box['segmentation']:
  117. gt_poly[i] = box['segmentation']
  118. has_segmentation = True
  119. if has_segmentation and not any(gt_poly):
  120. continue
  121. im_info = {
  122. 'im_id': np.array([img_id]).astype('int32'),
  123. 'image_shape': np.array([im_h, im_w]).astype('int32'),
  124. }
  125. label_info = {
  126. 'is_crowd': is_crowd,
  127. 'gt_class': gt_class,
  128. 'gt_bbox': gt_bbox,
  129. 'gt_score': gt_score,
  130. 'gt_poly': gt_poly,
  131. 'difficult': difficult
  132. }
  133. if None in gt_poly:
  134. del label_info['gt_poly']
  135. self.file_list.append(({
  136. 'image': im_fname,
  137. **
  138. im_info,
  139. **
  140. label_info
  141. }))
  142. if not len(self.file_list) > 0:
  143. raise Exception('not found any coco record in %s' % ann_file)
  144. logging.info("{} samples in file {}".format(
  145. len(self.file_list), ann_file))
  146. self.num_samples = len(self.file_list)
  147. self._epoch = 0