split_dataset.py 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116
  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 shutil
  13. import random
  14. import json
  15. from tqdm import tqdm
  16. from .....utils.file_interface import custom_open, write_json_file
  17. from .....utils.logging import info
  18. def split_dataset(root_dir, train_rate, val_rate):
  19. """ split dataset """
  20. assert train_rate + val_rate == 100, \
  21. f"The sum of train_rate({train_rate}), val_rate({val_rate}) should equal 100!"
  22. assert train_rate > 0 and val_rate > 0, \
  23. f"The train_rate({train_rate}) and val_rate({val_rate}) should be greater than 0!"
  24. all_image_info_list = []
  25. all_category_dict = {}
  26. max_image_id = 0
  27. for fn in ["instance_train.json", "instance_val.json"]:
  28. anno_path = os.path.join(root_dir, "annotations", fn)
  29. if not os.path.exists(anno_path):
  30. info(
  31. f"The annotation file {anno_path} don't exists, has been ignored!"
  32. )
  33. continue
  34. image_info_list, category_list, max_image_id = json2list(anno_path,
  35. max_image_id)
  36. all_image_info_list.extend(image_info_list)
  37. for category in category_list:
  38. if category['id'] not in all_category_dict:
  39. all_category_dict[category['id']] = category
  40. total_num = len(all_image_info_list)
  41. random.shuffle(all_image_info_list)
  42. all_category_list = [all_category_dict[k] for k in all_category_dict]
  43. start = 0
  44. for fn, rate in [("instance_train.json", train_rate),
  45. ("instance_val.json", val_rate)]:
  46. end = start + round(total_num * rate / 100)
  47. save_path = os.path.join(root_dir, "annotations", fn)
  48. if os.path.exists(save_path):
  49. bak_path = save_path + ".bak"
  50. shutil.move(save_path, bak_path)
  51. info(
  52. f"The original annotation file {fn} has been backed up to {bak_path}."
  53. )
  54. assemble_write(all_image_info_list[start:end], all_category_list,
  55. save_path)
  56. start = end
  57. return root_dir
  58. def json2list(json_path, base_image_num):
  59. """ load json as list """
  60. assert os.path.exists(json_path), json_path
  61. with custom_open(json_path, 'r') as f:
  62. data = json.load(f)
  63. image_info_dict = {}
  64. max_image_id = 0
  65. for image_info in data['images']:
  66. # 得到全局唯一的image_id
  67. global_image_id = image_info['id'] + base_image_num
  68. max_image_id = max(global_image_id, max_image_id)
  69. image_info['id'] = global_image_id
  70. image_info_dict[global_image_id] = {"img": image_info, 'anno': []}
  71. image_info_dict = {
  72. image_info['id']: {
  73. "img": image_info,
  74. 'anno': []
  75. }
  76. for image_info in data['images']
  77. }
  78. info(f"Start loading annotation file {json_path}...")
  79. for anno in tqdm(data['annotations']):
  80. global_image_id = anno['image_id'] + base_image_num
  81. anno['image_id'] = global_image_id
  82. image_info_dict[global_image_id]['anno'].append(anno)
  83. image_info_list = [(image_info_dict[image_info]['img'],
  84. image_info_dict[image_info]['anno'])
  85. for image_info in image_info_dict]
  86. return image_info_list, data['categories'], max_image_id
  87. def assemble_write(image_info_list, category_list, save_path):
  88. """ assemble coco format and save to file """
  89. coco_data = {'categories': category_list}
  90. image_list = [i[0] for i in image_info_list]
  91. all_anno_list = []
  92. for i in image_info_list:
  93. all_anno_list.extend(i[1])
  94. anno_list = []
  95. for i, anno in enumerate(all_anno_list):
  96. anno['id'] = i + 1
  97. anno_list.append(anno)
  98. coco_data['images'] = image_list
  99. coco_data['annotations'] = anno_list
  100. write_json_file(coco_data, save_path)
  101. info(f"The splited annotations has been save to {save_path}.")