mini_deep_globe_road_extraction.py 3.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  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. from .dataset import Dataset
  16. from paddlex.paddleseg.utils.download import download_file_and_uncompress
  17. from paddlex.paddleseg.utils import seg_env
  18. from paddlex.paddleseg.cvlibs import manager
  19. from paddlex.paddleseg.transforms import Compose
  20. URL = "https://paddleseg.bj.bcebos.com/dataset/MiniDeepGlobeRoadExtraction.zip"
  21. @manager.DATASETS.add_component
  22. class MiniDeepGlobeRoadExtraction(Dataset):
  23. """
  24. MiniDeepGlobeRoadExtraction dataset is extraced from DeepGlobe CVPR2018 challenge (http://deepglobe.org/)
  25. There are 800 images in the training set and 200 images in the validation set.
  26. Args:
  27. dataset_root (str, optional): The dataset directory. Default: None.
  28. transforms (list, optional): Transforms for image. Default: None.
  29. mode (str, optional): Which part of dataset to use. It is one of ('train', 'val'). Default: 'train'.
  30. edge (bool, optional): Whether to compute edge while training. Default: False.
  31. """
  32. NUM_CLASSES = 2
  33. def __init__(self,
  34. dataset_root=None,
  35. transforms=None,
  36. mode='train',
  37. edge=False):
  38. self.dataset_root = dataset_root
  39. self.transforms = Compose(transforms)
  40. mode = mode.lower()
  41. self.mode = mode
  42. self.file_list = list()
  43. self.num_classes = self.NUM_CLASSES
  44. self.ignore_index = 255
  45. self.edge = edge
  46. if mode not in ['train', 'val']:
  47. raise ValueError(
  48. "`mode` should be 'train' or 'val', but got {}.".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. 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. file_path = os.path.join(self.dataset_root, 'train.txt')
  67. else:
  68. file_path = os.path.join(self.dataset_root, 'val.txt')
  69. with open(file_path, 'r') as f:
  70. for line in f:
  71. items = line.strip().split('|')
  72. if len(items) != 2:
  73. if mode == 'train' or mode == 'val':
  74. raise Exception(
  75. "File list format incorrect! It should be"
  76. " image_name|label_name\\n")
  77. image_path = os.path.join(self.dataset_root, items[0])
  78. grt_path = None
  79. else:
  80. image_path = os.path.join(self.dataset_root, items[0])
  81. grt_path = os.path.join(self.dataset_root, items[1])
  82. self.file_list.append([image_path, grt_path])