dataset.py 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192
  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. from collections import OrderedDict
  17. try:
  18. from collections.abc import Sequence
  19. except Exception:
  20. from collections import Sequence
  21. from paddle.io import Dataset
  22. from paddlex.ppdet.core.workspace import register, serializable
  23. from paddlex.ppdet.utils.download import get_dataset_path
  24. import copy
  25. @serializable
  26. class DetDataset(Dataset):
  27. """
  28. Load detection dataset.
  29. Args:
  30. dataset_dir (str): root directory for dataset.
  31. image_dir (str): directory for images.
  32. anno_path (str): annotation file path.
  33. data_fields (list): key name of data dictionary, at least have 'image'.
  34. sample_num (int): number of samples to load, -1 means all.
  35. use_default_label (bool): whether to load default label list.
  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. use_default_label=None,
  44. **kwargs):
  45. super(DetDataset, self).__init__()
  46. self.dataset_dir = dataset_dir if dataset_dir is not None else ''
  47. self.anno_path = anno_path
  48. self.image_dir = image_dir if image_dir is not None else ''
  49. self.data_fields = data_fields
  50. self.sample_num = sample_num
  51. self.use_default_label = use_default_label
  52. self._epoch = 0
  53. self._curr_iter = 0
  54. def __len__(self, ):
  55. return len(self.roidbs)
  56. def __getitem__(self, idx):
  57. # data batch
  58. roidb = copy.deepcopy(self.roidbs[idx])
  59. if self.mixup_epoch == 0 or self._epoch < self.mixup_epoch:
  60. n = len(self.roidbs)
  61. idx = np.random.randint(n)
  62. roidb = [roidb, copy.deepcopy(self.roidbs[idx])]
  63. elif self.cutmix_epoch == 0 or self._epoch < self.cutmix_epoch:
  64. n = len(self.roidbs)
  65. idx = np.random.randint(n)
  66. roidb = [roidb, copy.deepcopy(self.roidbs[idx])]
  67. elif self.mosaic_epoch == 0 or self._epoch < self.mosaic_epoch:
  68. n = len(self.roidbs)
  69. roidb = [roidb, ] + [
  70. copy.deepcopy(self.roidbs[np.random.randint(n)])
  71. for _ in range(3)
  72. ]
  73. if isinstance(roidb, Sequence):
  74. for r in roidb:
  75. r['curr_iter'] = self._curr_iter
  76. else:
  77. roidb['curr_iter'] = self._curr_iter
  78. self._curr_iter += 1
  79. return self.transform(roidb)
  80. def check_or_download_dataset(self):
  81. self.dataset_dir = get_dataset_path(self.dataset_dir, self.anno_path,
  82. self.image_dir)
  83. def set_kwargs(self, **kwargs):
  84. self.mixup_epoch = kwargs.get('mixup_epoch', -1)
  85. self.cutmix_epoch = kwargs.get('cutmix_epoch', -1)
  86. self.mosaic_epoch = kwargs.get('mosaic_epoch', -1)
  87. def set_transform(self, transform):
  88. self.transform = transform
  89. def set_epoch(self, epoch_id):
  90. self._epoch = epoch_id
  91. def parse_dataset(self, ):
  92. raise NotImplementedError(
  93. "Need to implement parse_dataset method of Dataset")
  94. def get_anno(self):
  95. if self.anno_path is None:
  96. return
  97. return os.path.join(self.dataset_dir, self.anno_path)
  98. def _is_valid_file(f, extensions=('.jpg', '.jpeg', '.png', '.bmp')):
  99. return f.lower().endswith(extensions)
  100. def _make_dataset(dir):
  101. dir = os.path.expanduser(dir)
  102. if not os.path.isdir(dir):
  103. raise ('{} should be a dir'.format(dir))
  104. images = []
  105. for root, _, fnames in sorted(os.walk(dir, followlinks=True)):
  106. for fname in sorted(fnames):
  107. path = os.path.join(root, fname)
  108. if _is_valid_file(path):
  109. images.append(path)
  110. return images
  111. @register
  112. @serializable
  113. class ImageFolder(DetDataset):
  114. def __init__(self,
  115. dataset_dir=None,
  116. image_dir=None,
  117. anno_path=None,
  118. sample_num=-1,
  119. use_default_label=None,
  120. **kwargs):
  121. super(ImageFolder, self).__init__(
  122. dataset_dir,
  123. image_dir,
  124. anno_path,
  125. sample_num=sample_num,
  126. use_default_label=use_default_label)
  127. self._imid2path = {}
  128. self.roidbs = None
  129. self.sample_num = sample_num
  130. def check_or_download_dataset(self):
  131. return
  132. def parse_dataset(self, ):
  133. if not self.roidbs:
  134. self.roidbs = self._load_images()
  135. def _parse(self):
  136. image_dir = self.image_dir
  137. if not isinstance(image_dir, Sequence):
  138. image_dir = [image_dir]
  139. images = []
  140. for im_dir in image_dir:
  141. if os.path.isdir(im_dir):
  142. im_dir = os.path.join(self.dataset_dir, im_dir)
  143. images.extend(_make_dataset(im_dir))
  144. elif os.path.isfile(im_dir) and _is_valid_file(im_dir):
  145. images.append(im_dir)
  146. return images
  147. def _load_images(self):
  148. images = self._parse()
  149. ct = 0
  150. records = []
  151. for image in images:
  152. assert image != '' and os.path.isfile(image), \
  153. "Image {} not found".format(image)
  154. if self.sample_num > 0 and ct >= self.sample_num:
  155. break
  156. rec = {'im_id': np.array([ct]), 'im_file': image}
  157. self._imid2path[ct] = image
  158. ct += 1
  159. records.append(rec)
  160. assert len(records) > 0, "No image file found"
  161. return records
  162. def get_imid2path(self):
  163. return self._imid2path
  164. def set_images(self, images):
  165. self.image_dir = images
  166. self.roidbs = self._load_images()