ade.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  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 numpy as np
  16. from PIL import Image
  17. from paddlex.paddleseg.datasets import Dataset
  18. from paddlex.paddleseg.utils.download import download_file_and_uncompress
  19. from paddlex.paddleseg.utils import seg_env
  20. from paddlex.paddleseg.cvlibs import manager
  21. from paddlex.paddleseg.transforms import Compose
  22. import paddlex.paddleseg.transforms.functional as F
  23. URL = "http://data.csail.mit.edu/places/ADEchallenge/ADEChallengeData2016.zip"
  24. @manager.DATASETS.add_component
  25. class ADE20K(Dataset):
  26. """
  27. ADE20K dataset `http://sceneparsing.csail.mit.edu/`.
  28. Args:
  29. transforms (list): A list of image transformations.
  30. dataset_root (str, optional): The ADK20K dataset directory. Default: None.
  31. mode (str, optional): A subset of the entire dataset. It should be one of ('train', 'val'). Default: 'train'.
  32. edge (bool, optional): Whether to compute edge while training. Default: False
  33. """
  34. NUM_CLASSES = 150
  35. def __init__(self, transforms, dataset_root=None, mode='train', edge=False):
  36. self.dataset_root = dataset_root
  37. self.transforms = Compose(transforms)
  38. mode = mode.lower()
  39. self.mode = mode
  40. self.file_list = list()
  41. self.num_classes = self.NUM_CLASSES
  42. self.ignore_index = 255
  43. self.edge = edge
  44. if mode not in ['train', 'val']:
  45. raise ValueError(
  46. "`mode` should be one of ('train', 'val') in ADE20K dataset, but got {}."
  47. .format(mode))
  48. if self.transforms is None:
  49. raise ValueError("`transforms` is necessary, but it is None.")
  50. if self.dataset_root is None:
  51. self.dataset_root = download_file_and_uncompress(
  52. url=URL,
  53. savepath=seg_env.DATA_HOME,
  54. extrapath=seg_env.DATA_HOME,
  55. extraname='ADEChallengeData2016')
  56. elif not os.path.exists(self.dataset_root):
  57. self.dataset_root = os.path.normpath(self.dataset_root)
  58. savepath, extraname = self.dataset_root.rsplit(
  59. sep=os.path.sep, maxsplit=1)
  60. self.dataset_root = download_file_and_uncompress(
  61. url=URL,
  62. savepath=savepath,
  63. extrapath=savepath,
  64. extraname=extraname)
  65. if mode == 'train':
  66. img_dir = os.path.join(self.dataset_root, 'images/training')
  67. label_dir = os.path.join(self.dataset_root, 'annotations/training')
  68. elif mode == 'val':
  69. img_dir = os.path.join(self.dataset_root, 'images/validation')
  70. label_dir = os.path.join(self.dataset_root,
  71. 'annotations/validation')
  72. img_files = os.listdir(img_dir)
  73. label_files = [i.replace('.jpg', '.png') for i in img_files]
  74. for i in range(len(img_files)):
  75. img_path = os.path.join(img_dir, img_files[i])
  76. label_path = os.path.join(label_dir, label_files[i])
  77. self.file_list.append([img_path, label_path])
  78. def __getitem__(self, idx):
  79. image_path, label_path = self.file_list[idx]
  80. if self.mode == 'val':
  81. im, _ = self.transforms(im=image_path)
  82. label = np.asarray(Image.open(label_path))
  83. # The class 0 is ignored. And it will equal to 255 after
  84. # subtracted 1, because the dtype of label is uint8.
  85. label = label - 1
  86. label = label[np.newaxis, :, :]
  87. return im, label
  88. else:
  89. im, label = self.transforms(im=image_path, label=label_path)
  90. label = label - 1
  91. # Recover the ignore pixels adding by transform
  92. label[label == 254] = 255
  93. if self.edge:
  94. edge_mask = F.mask_to_binary_edge(
  95. label, radius=2, num_classes=self.num_classes)
  96. return im, label, edge_mask
  97. else:
  98. return im, label