split_dataset.py 2.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. # !/usr/bin/env python3
  2. # -*- coding: UTF-8 -*-
  3. ################################################################################
  4. #
  5. # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
  6. #
  7. ################################################################################
  8. """
  9. Author: PaddlePaddle Authors
  10. """
  11. import os
  12. import os.path as osp
  13. import random
  14. import shutil
  15. from .....utils.file_interface import custom_open
  16. from .....utils import logging
  17. def split_dataset(root_dir, train_percent, val_percent):
  18. """ split dataset """
  19. assert train_percent > 0, ValueError(
  20. f"The train_percent({train_percent}) must greater than 0!")
  21. assert val_percent > 0, ValueError(
  22. f"The val_percent({val_percent}) must greater than 0!")
  23. if train_percent + val_percent != 100:
  24. raise ValueError(
  25. f"The sum of train_percent({train_percent})and val_percent({val_percent}) should be 100!"
  26. )
  27. img_dir = osp.join(root_dir, "images")
  28. assert osp.exists(img_dir), FileNotFoundError(
  29. f"The dir of images ({img_dir}) doesn't exist, please check!")
  30. ann_dir = osp.join(root_dir, "annotations")
  31. assert osp.exists(ann_dir), FileNotFoundError(
  32. f"The dir of annotations ({ann_dir}) doesn't exist, please check!")
  33. img_file_list = [
  34. osp.join("images", img_name) for img_name in os.listdir(img_dir)
  35. ]
  36. img_num = len(img_file_list)
  37. ann_file_list = [
  38. osp.join("annotations", ann_name) for ann_name in os.listdir(ann_dir)
  39. ]
  40. ann_num = len(ann_file_list)
  41. assert img_num == ann_num, ValueError(
  42. "The number of images and annotations must be equal!")
  43. split_tags = ["train", "val"]
  44. mapping_line_list = []
  45. for tag in split_tags:
  46. mapping_file = osp.join(root_dir, f"{tag}.txt")
  47. if not osp.exists(mapping_file):
  48. logging.info(
  49. f"The mapping file ({mapping_file}) doesn't exist, ignored.")
  50. continue
  51. with custom_open(mapping_file, "r") as fp:
  52. lines = filter(None, (line.strip() for line in fp.readlines()))
  53. mapping_line_list.extend(lines)
  54. sample_num = len(mapping_line_list)
  55. random.shuffle(mapping_line_list)
  56. split_percents = [train_percent, val_percent]
  57. start_idx = 0
  58. for tag, percent in zip(split_tags, split_percents):
  59. if tag == 'test' and percent == 0:
  60. continue
  61. end_idx = start_idx + round(sample_num * percent / 100)
  62. end_idx = min(end_idx, sample_num)
  63. mapping_file = osp.join(root_dir, f"{tag}.txt")
  64. if os.path.exists(mapping_file):
  65. shutil.move(mapping_file, mapping_file + ".bak")
  66. logging.info(f"The original mapping file ({mapping_file}) "
  67. f"has been backed up to ({mapping_file}.bak)")
  68. with custom_open(mapping_file, "w") as fp:
  69. fp.write("\n".join(mapping_line_list[start_idx:end_idx]))
  70. start_idx = end_idx
  71. return root_dir