ade.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112
  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. from paddlex.paddleseg.transforms import 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',
  36. edge=False):
  37. self.dataset_root = dataset_root
  38. self.transforms = Compose(transforms)
  39. mode = mode.lower()
  40. self.mode = mode
  41. self.file_list = list()
  42. self.num_classes = self.NUM_CLASSES
  43. self.ignore_index = 255
  44. self.edge = edge
  45. if mode not in ['train', 'val']:
  46. raise ValueError(
  47. "`mode` should be one of ('train', 'val') in ADE20K dataset, but got {}."
  48. .format(mode))
  49. if self.transforms is None:
  50. raise ValueError("`transforms` is necessary, but it is None.")
  51. if self.dataset_root is None:
  52. self.dataset_root = download_file_and_uncompress(
  53. url=URL,
  54. savepath=seg_env.DATA_HOME,
  55. extrapath=seg_env.DATA_HOME,
  56. extraname='ADEChallengeData2016')
  57. elif not os.path.exists(self.dataset_root):
  58. self.dataset_root = os.path.normpath(self.dataset_root)
  59. savepath, extraname = self.dataset_root.rsplit(
  60. sep=os.path.sep, maxsplit=1)
  61. self.dataset_root = download_file_and_uncompress(
  62. url=URL,
  63. savepath=savepath,
  64. extrapath=savepath,
  65. extraname=extraname)
  66. if mode == 'train':
  67. img_dir = os.path.join(self.dataset_root, 'images/training')
  68. label_dir = os.path.join(self.dataset_root, 'annotations/training')
  69. elif mode == 'val':
  70. img_dir = os.path.join(self.dataset_root, 'images/validation')
  71. label_dir = os.path.join(self.dataset_root,
  72. 'annotations/validation')
  73. img_files = os.listdir(img_dir)
  74. label_files = [i.replace('.jpg', '.png') for i in img_files]
  75. for i in range(len(img_files)):
  76. img_path = os.path.join(img_dir, img_files[i])
  77. label_path = os.path.join(label_dir, label_files[i])
  78. self.file_list.append([img_path, label_path])
  79. def __getitem__(self, idx):
  80. image_path, label_path = self.file_list[idx]
  81. if self.mode == 'val':
  82. im, _ = self.transforms(im=image_path)
  83. label = np.asarray(Image.open(label_path))
  84. # The class 0 is ignored. And it will equal to 255 after
  85. # subtracted 1, because the dtype of label is uint8.
  86. label = label - 1
  87. label = label[np.newaxis, :, :]
  88. return im, label
  89. else:
  90. im, label = self.transforms(im=image_path, label=label_path)
  91. label = label - 1
  92. # Recover the ignore pixels adding by transform
  93. label[label == 254] = 255
  94. if self.edge:
  95. edge_mask = F.mask_to_binary_edge(
  96. label, radius=2, num_classes=self.num_classes)
  97. return im, label, edge_mask
  98. else:
  99. return im, label