prepare_levircd.py 1.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142
  1. #!/usr/bin/env python
  2. import os.path as osp
  3. from common import (get_default_parser, add_crop_options, crop_patches,
  4. get_path_tuples, create_file_list, link_dataset)
  5. SUBSETS = ('train', 'val', 'test')
  6. SUBDIRS = ('A', 'B', 'label')
  7. FILE_LIST_PATTERN = "{subset}.txt"
  8. URL = ""
  9. if __name__ == '__main__':
  10. parser = get_default_parser()
  11. parser = add_crop_options(parser)
  12. args = parser.parse_args()
  13. out_dir = osp.join(args.out_dataset_dir,
  14. osp.basename(osp.normpath(args.in_dataset_dir)))
  15. if args.crop_size is not None:
  16. crop_patches(
  17. args.crop_size,
  18. args.crop_stride,
  19. data_dir=args.in_dataset_dir,
  20. out_dir=out_dir,
  21. subsets=SUBSETS,
  22. subdirs=SUBDIRS,
  23. glob_pattern='*.png',
  24. max_workers=0)
  25. else:
  26. link_dataset(args.in_dataset_dir, args.out_dataset_dir)
  27. for subset in SUBSETS:
  28. path_tuples = get_path_tuples(
  29. *(osp.join(out_dir, subset, subdir) for subdir in SUBDIRS),
  30. glob_pattern='**/*.png',
  31. data_dir=args.out_dataset_dir)
  32. file_list = osp.join(
  33. args.out_dataset_dir, FILE_LIST_PATTERN.format(subset=subset))
  34. create_file_list(file_list, path_tuples)
  35. print(f"Write file list to {file_list}.")