dataset.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164
  1. # Copyright (c) 2020 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 paddle
  16. import numpy as np
  17. from PIL import Image
  18. from paddlex.paddleseg.cvlibs import manager
  19. from paddlex.paddleseg.transforms import Compose
  20. from paddlex.paddleseg.transforms import functional as F
  21. @manager.DATASETS.add_component
  22. class Dataset(paddle.io.Dataset):
  23. """
  24. Pass in a custom dataset that conforms to the format.
  25. Args:
  26. transforms (list): Transforms for image.
  27. dataset_root (str): The dataset directory.
  28. num_classes (int): Number of classes.
  29. mode (str, optional): which part of dataset to use. it is one of ('train', 'val', 'test'). Default: 'train'.
  30. train_path (str, optional): The train dataset file. When mode is 'train', train_path is necessary.
  31. The contents of train_path file are as follow:
  32. image1.jpg ground_truth1.png
  33. image2.jpg ground_truth2.png
  34. val_path (str. optional): The evaluation dataset file. When mode is 'val', val_path is necessary.
  35. The contents is the same as train_path
  36. test_path (str, optional): The test dataset file. When mode is 'test', test_path is necessary.
  37. The annotation file is not necessary in test_path file.
  38. separator (str, optional): The separator of dataset list. Default: ' '.
  39. edge (bool, optional): Whether to compute edge while training. Default: False
  40. Examples:
  41. import paddlex.paddleseg.transforms as T
  42. from paddlex.paddleseg.datasets import Dataset
  43. transforms = [T.RandomPaddingCrop(crop_size=(512,512)), T.Normalize()]
  44. dataset_root = 'dataset_root_path'
  45. train_path = 'train_path'
  46. num_classes = 2
  47. dataset = Dataset(transforms = transforms,
  48. dataset_root = dataset_root,
  49. num_classes = 2,
  50. train_path = train_path,
  51. mode = 'train')
  52. """
  53. def __init__(self,
  54. transforms,
  55. dataset_root,
  56. num_classes,
  57. mode='train',
  58. train_path=None,
  59. val_path=None,
  60. test_path=None,
  61. separator=' ',
  62. ignore_index=255,
  63. edge=False):
  64. self.dataset_root = dataset_root
  65. self.transforms = Compose(transforms)
  66. self.file_list = list()
  67. mode = mode.lower()
  68. self.mode = mode
  69. self.num_classes = num_classes
  70. self.ignore_index = ignore_index
  71. self.edge = edge
  72. if mode.lower() not in ['train', 'val', 'test']:
  73. raise ValueError(
  74. "mode should be 'train', 'val' or 'test', but got {}.".format(
  75. mode))
  76. if self.transforms is None:
  77. raise ValueError("`transforms` is necessary, but it is None.")
  78. self.dataset_root = dataset_root
  79. if not os.path.exists(self.dataset_root):
  80. raise FileNotFoundError('there is not `dataset_root`: {}.'.format(
  81. self.dataset_root))
  82. if mode == 'train':
  83. if train_path is None:
  84. raise ValueError(
  85. 'When `mode` is "train", `train_path` is necessary, but it is None.'
  86. )
  87. elif not os.path.exists(train_path):
  88. raise FileNotFoundError('`train_path` is not found: {}'.format(
  89. train_path))
  90. else:
  91. file_path = train_path
  92. elif mode == 'val':
  93. if val_path is None:
  94. raise ValueError(
  95. 'When `mode` is "val", `val_path` is necessary, but it is None.'
  96. )
  97. elif not os.path.exists(val_path):
  98. raise FileNotFoundError('`val_path` is not found: {}'.format(
  99. val_path))
  100. else:
  101. file_path = val_path
  102. else:
  103. if test_path is None:
  104. raise ValueError(
  105. 'When `mode` is "test", `test_path` is necessary, but it is None.'
  106. )
  107. elif not os.path.exists(test_path):
  108. raise FileNotFoundError('`test_path` is not found: {}'.format(
  109. test_path))
  110. else:
  111. file_path = test_path
  112. with open(file_path, 'r') as f:
  113. for line in f:
  114. items = line.strip().split(separator)
  115. if len(items) != 2:
  116. if mode == 'train' or mode == 'val':
  117. raise ValueError(
  118. "File list format incorrect! In training or evaluation task it should be"
  119. " image_name{}label_name\\n".format(separator))
  120. image_path = os.path.join(self.dataset_root, items[0])
  121. label_path = None
  122. else:
  123. image_path = os.path.join(self.dataset_root, items[0])
  124. label_path = os.path.join(self.dataset_root, items[1])
  125. self.file_list.append([image_path, label_path])
  126. def __getitem__(self, idx):
  127. image_path, label_path = self.file_list[idx]
  128. if self.mode == 'test':
  129. im, _ = self.transforms(im=image_path)
  130. im = im[np.newaxis, ...]
  131. return im, image_path
  132. elif self.mode == 'val':
  133. im, _ = self.transforms(im=image_path)
  134. label = np.asarray(Image.open(label_path))
  135. label = label[np.newaxis, :, :]
  136. return im, label
  137. else:
  138. im, label = self.transforms(im=image_path, label=label_path)
  139. if self.edge:
  140. edge_mask = F.mask_to_binary_edge(
  141. label, radius=2, num_classes=self.num_classes)
  142. return im, label, edge_mask
  143. else:
  144. return im, label
  145. def __len__(self):
  146. return len(self.file_list)