sniper_coco.py 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191
  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 json
  17. import copy
  18. import numpy as np
  19. try:
  20. from collections.abc import Sequence
  21. except Exception:
  22. from collections import Sequence
  23. from paddlex.ppdet.core.workspace import register, serializable
  24. from paddlex.ppdet.data.crop_utils.annotation_cropper import AnnoCropper
  25. from .coco import COCODataSet
  26. from .dataset import _make_dataset, _is_valid_file
  27. from paddlex.ppdet.utils.logger import setup_logger
  28. logger = setup_logger('sniper_coco_dataset')
  29. @register
  30. @serializable
  31. class SniperCOCODataSet(COCODataSet):
  32. """SniperCOCODataSet"""
  33. def __init__(self,
  34. dataset_dir=None,
  35. image_dir=None,
  36. anno_path=None,
  37. proposals_file=None,
  38. data_fields=['image'],
  39. sample_num=-1,
  40. load_crowd=False,
  41. allow_empty=True,
  42. empty_ratio=1.,
  43. is_trainset=True,
  44. image_target_sizes=[2000, 1000],
  45. valid_box_ratio_ranges=[[-1, 0.1], [0.08, -1]],
  46. chip_target_size=500,
  47. chip_target_stride=200,
  48. use_neg_chip=False,
  49. max_neg_num_per_im=8,
  50. max_per_img=-1,
  51. nms_thresh=0.5):
  52. super(SniperCOCODataSet, self).__init__(
  53. dataset_dir=dataset_dir,
  54. image_dir=image_dir,
  55. anno_path=anno_path,
  56. data_fields=data_fields,
  57. sample_num=sample_num,
  58. load_crowd=load_crowd,
  59. allow_empty=allow_empty,
  60. empty_ratio=empty_ratio)
  61. self.proposals_file = proposals_file
  62. self.proposals = None
  63. self.anno_cropper = None
  64. self.is_trainset = is_trainset
  65. self.image_target_sizes = image_target_sizes
  66. self.valid_box_ratio_ranges = valid_box_ratio_ranges
  67. self.chip_target_size = chip_target_size
  68. self.chip_target_stride = chip_target_stride
  69. self.use_neg_chip = use_neg_chip
  70. self.max_neg_num_per_im = max_neg_num_per_im
  71. self.max_per_img = max_per_img
  72. self.nms_thresh = nms_thresh
  73. def parse_dataset(self):
  74. if not hasattr(self, "roidbs"):
  75. super(SniperCOCODataSet, self).parse_dataset()
  76. if self.is_trainset:
  77. self._parse_proposals()
  78. self._merge_anno_proposals()
  79. self.ori_roidbs = copy.deepcopy(self.roidbs)
  80. self.init_anno_cropper()
  81. self.roidbs = self.generate_chips_roidbs(self.roidbs, self.is_trainset)
  82. def set_proposals_file(self, file_path):
  83. self.proposals_file = file_path
  84. def init_anno_cropper(self):
  85. logger.info("Init AnnoCropper...")
  86. self.anno_cropper = AnnoCropper(
  87. image_target_sizes=self.image_target_sizes,
  88. valid_box_ratio_ranges=self.valid_box_ratio_ranges,
  89. chip_target_size=self.chip_target_size,
  90. chip_target_stride=self.chip_target_stride,
  91. use_neg_chip=self.use_neg_chip,
  92. max_neg_num_per_im=self.max_neg_num_per_im,
  93. max_per_img=self.max_per_img,
  94. nms_thresh=self.nms_thresh)
  95. def generate_chips_roidbs(self, roidbs, is_trainset):
  96. if is_trainset:
  97. roidbs = self.anno_cropper.crop_anno_records(roidbs)
  98. else:
  99. roidbs = self.anno_cropper.crop_infer_anno_records(roidbs)
  100. return roidbs
  101. def _parse_proposals(self):
  102. if self.proposals_file:
  103. self.proposals = {}
  104. logger.info("Parse proposals file:{}".format(self.proposals_file))
  105. with open(self.proposals_file, 'r') as f:
  106. proposals = json.load(f)
  107. for prop in proposals:
  108. image_id = prop["image_id"]
  109. if image_id not in self.proposals:
  110. self.proposals[image_id] = []
  111. x, y, w, h = prop["bbox"]
  112. self.proposals[image_id].append([x, y, x + w, y + h])
  113. def _merge_anno_proposals(self):
  114. assert self.roidbs
  115. if self.proposals and len(self.proposals.keys()) > 0:
  116. logger.info("merge proposals to annos")
  117. for id, record in enumerate(self.roidbs):
  118. image_id = int(record["im_id"])
  119. if image_id not in self.proposals.keys():
  120. logger.info("image id :{} no proposals".format(image_id))
  121. record["proposals"] = np.array(
  122. self.proposals.get(image_id, []), dtype=np.float32)
  123. self.roidbs[id] = record
  124. def get_ori_roidbs(self):
  125. if not hasattr(self, "ori_roidbs"):
  126. return None
  127. return self.ori_roidbs
  128. def get_roidbs(self):
  129. if not hasattr(self, "roidbs"):
  130. self.parse_dataset()
  131. return self.roidbs
  132. def set_roidbs(self, roidbs):
  133. self.roidbs = roidbs
  134. def check_or_download_dataset(self):
  135. return
  136. def _parse(self):
  137. image_dir = self.image_dir
  138. if not isinstance(image_dir, Sequence):
  139. image_dir = [image_dir]
  140. images = []
  141. for im_dir in image_dir:
  142. if os.path.isdir(im_dir):
  143. im_dir = os.path.join(self.dataset_dir, im_dir)
  144. images.extend(_make_dataset(im_dir))
  145. elif os.path.isfile(im_dir) and _is_valid_file(im_dir):
  146. images.append(im_dir)
  147. return images
  148. def _load_images(self):
  149. images = self._parse()
  150. ct = 0
  151. records = []
  152. for image in images:
  153. assert image != '' and os.path.isfile(image), \
  154. "Image {} not found".format(image)
  155. if self.sample_num > 0 and ct >= self.sample_num:
  156. break
  157. im = cv2.imread(image)
  158. h, w, c = im.shape
  159. rec = {'im_id': np.array([ct]), 'im_file': image, "h": h, "w": w}
  160. self._imid2path[ct] = image
  161. ct += 1
  162. records.append(rec)
  163. assert len(records) > 0, "No image file found"
  164. return records
  165. def get_imid2path(self):
  166. return self._imid2path
  167. def set_images(self, images):
  168. self._imid2path = {}
  169. self.image_dir = images
  170. self.roidbs = self._load_images()