coco_split.py 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  1. # Copyright (c) 2021 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.path as osp
  15. import random
  16. import json
  17. from paddlex.utils import MyEncoder, logging
  18. def split_coco_dataset(dataset_dir, val_percent, test_percent, save_dir):
  19. # matplotlib.use() must be called *before* pylab, matplotlib.pyplot,
  20. # or matplotlib.backends is imported for the first time
  21. # pycocotools import matplotlib
  22. import matplotlib
  23. matplotlib.use('Agg')
  24. from pycocotools.coco import COCO
  25. if not osp.exists(osp.join(dataset_dir, "annotations.json")):
  26. logging.error("\'annotations.json\' is not found in {}!".format(
  27. dataset_dir))
  28. annotation_file = osp.join(dataset_dir, "annotations.json")
  29. coco = COCO(annotation_file)
  30. img_ids = coco.getImgIds()
  31. cat_ids = coco.getCatIds()
  32. anno_ids = coco.getAnnIds()
  33. val_num = int(len(img_ids) * val_percent)
  34. test_num = int(len(img_ids) * test_percent)
  35. train_num = len(img_ids) - val_num - test_num
  36. random.shuffle(img_ids)
  37. train_files_ids = img_ids[:train_num]
  38. val_files_ids = img_ids[train_num:train_num + val_num]
  39. test_files_ids = img_ids[train_num + val_num:]
  40. for img_id_list in [train_files_ids, val_files_ids, test_files_ids]:
  41. img_anno_ids = coco.getAnnIds(imgIds=img_id_list, iscrowd=0)
  42. imgs = coco.loadImgs(img_id_list)
  43. instances = coco.loadAnns(img_anno_ids)
  44. categories = coco.loadCats(cat_ids)
  45. img_dict = {
  46. "annotations": instances,
  47. "images": imgs,
  48. "categories": categories
  49. }
  50. if img_id_list == train_files_ids:
  51. json_file = open(osp.join(save_dir, 'train.json'), 'w+')
  52. json.dump(img_dict, json_file, cls=MyEncoder)
  53. elif img_id_list == val_files_ids:
  54. json_file = open(osp.join(save_dir, 'val.json'), 'w+')
  55. json.dump(img_dict, json_file, cls=MyEncoder)
  56. elif img_id_list == test_files_ids and len(test_files_ids):
  57. json_file = open(osp.join(save_dir, 'test.json'), 'w+')
  58. json.dump(img_dict, json_file, cls=MyEncoder)
  59. return train_num, val_num, test_num