prepare_ucmerced.py 2.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576
  1. #!/usr/bin/env python
  2. # Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. import random
  16. import os.path as osp
  17. from glob import iglob
  18. from functools import reduce, partial
  19. from common import (get_default_parser, create_file_list, link_dataset,
  20. random_split, create_label_list)
  21. CLASSES = ('agricultural', 'airplane', 'baseballdiamond', 'beach', 'buildings',
  22. 'chaparral', 'denseresidential', 'forest', 'freeway', 'golfcourse',
  23. 'harbor', 'intersection', 'mediumresidential', 'mobilehomepark',
  24. 'overpass', 'parkinglot', 'river', 'runway', 'sparseresidential',
  25. 'storagetanks', 'tenniscourt')
  26. SUBSETS = ('train', 'val', 'test')
  27. SUBDIRS = tuple(osp.join('Images', cls) for cls in CLASSES)
  28. FILE_LIST_PATTERN = "{subset}.txt"
  29. LABEL_LIST_NAME = "labels.txt"
  30. URL = ""
  31. if __name__ == '__main__':
  32. parser = get_default_parser()
  33. parser.add_argument('--seed', type=int, default=None, help="Random seed.")
  34. parser.add_argument(
  35. '--ratios',
  36. type=float,
  37. nargs='+',
  38. default=(0.7, 0.2, 0.1),
  39. help="Ratios of each subset (train/val or train/val/test).")
  40. args = parser.parse_args()
  41. if args.seed is not None:
  42. random.seed(args.seed)
  43. if len(args.ratios) not in (2, 3):
  44. raise ValueError("Wrong number of ratios!")
  45. out_dir = osp.join(args.out_dataset_dir,
  46. osp.basename(osp.normpath(args.in_dataset_dir)))
  47. link_dataset(args.in_dataset_dir, args.out_dataset_dir)
  48. splits_list = []
  49. for idx, (cls, subdir) in enumerate(zip(CLASSES, SUBDIRS)):
  50. pairs = []
  51. for p in iglob(osp.join(out_dir, subdir, '*.tif')):
  52. pair = (osp.relpath(p, args.out_dataset_dir), str(idx))
  53. pairs.append(pair)
  54. splits = random_split(pairs, ratios=args.ratios)
  55. splits_list.append(splits)
  56. splits = map(partial(reduce, list.__add__), zip(*splits_list))
  57. for subset, split in zip(SUBSETS, splits):
  58. file_list = osp.join(
  59. args.out_dataset_dir, FILE_LIST_PATTERN.format(subset=subset))
  60. create_file_list(file_list, split)
  61. print(f"Write file list to {file_list}.")
  62. label_list = osp.join(args.out_dataset_dir, LABEL_LIST_NAME)
  63. create_label_list(label_list, CLASSES)
  64. print(f"Write label list to {label_list}.")