optic_disc_seg.py 3.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697
  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/optic_disc_seg.zip"
  21. @manager.DATASETS.add_component
  22. class OpticDiscSeg(Dataset):
  23. """
  24. OpticDiscSeg dataset is extraced from iChallenge-AMD
  25. (https://ai.baidu.com/broad/subordinate?dataset=amd).
  26. Args:
  27. transforms (list): Transforms for image.
  28. dataset_root (str): The dataset directory. Default: None
  29. mode (str, optional): Which part of dataset to use. it is one of ('train', 'val', 'test'). 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', 'test']:
  47. raise ValueError(
  48. "`mode` should be 'train', 'val' or 'test', but got {}.".format(
  49. mode))
  50. if self.transforms is None:
  51. raise ValueError("`transforms` is necessary, but it is None.")
  52. if self.dataset_root is None:
  53. self.dataset_root = download_file_and_uncompress(
  54. url=URL,
  55. savepath=seg_env.DATA_HOME,
  56. extrapath=seg_env.DATA_HOME)
  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. file_path = os.path.join(self.dataset_root, 'train_list.txt')
  68. elif mode == 'val':
  69. file_path = os.path.join(self.dataset_root, 'val_list.txt')
  70. else:
  71. file_path = os.path.join(self.dataset_root, 'test_list.txt')
  72. with open(file_path, 'r') as f:
  73. for line in f:
  74. items = line.strip().split()
  75. if len(items) != 2:
  76. if mode == 'train' or mode == 'val':
  77. raise Exception(
  78. "File list format incorrect! It should be"
  79. " image_name label_name\\n")
  80. image_path = os.path.join(self.dataset_root, items[0])
  81. grt_path = None
  82. else:
  83. image_path = os.path.join(self.dataset_root, items[0])
  84. grt_path = os.path.join(self.dataset_root, items[1])
  85. self.file_list.append([image_path, grt_path])