convert_dataset.py 7.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225
  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 json
  14. import random
  15. import xml.etree.ElementTree as ET
  16. import cv2
  17. import numpy as np
  18. from PIL import Image, ImageDraw
  19. from tqdm import tqdm
  20. from .....utils.file_interface import custom_open, write_json_file
  21. from .....utils.errors import ConvertFailedError
  22. from .....utils.logging import info, warning
  23. class Indexer(object):
  24. """ Indexer """
  25. def __init__(self):
  26. """ init indexer """
  27. self._map = {}
  28. self.idx = 0
  29. def get_id(self, key):
  30. """ get id by key """
  31. if key not in self._map:
  32. self.idx += 1
  33. self._map[key] = self.idx
  34. return self._map[key]
  35. def get_list(self, key_name):
  36. """ return list containing key and id """
  37. map_list = []
  38. for key in self._map:
  39. val = self._map[key]
  40. map_list.append({key_name: key, 'id': val})
  41. return map_list
  42. class Extension(object):
  43. """ Extension """
  44. def __init__(self, exts_list):
  45. """ init extension """
  46. self._exts_list = ['.' + ext for ext in exts_list]
  47. def __iter__(self):
  48. """ iterator """
  49. return iter(self._exts_list)
  50. def update(self, ext):
  51. """ update extension """
  52. self._exts_list.remove(ext)
  53. self._exts_list.insert(0, ext)
  54. def check_src_dataset(root_dir, dataset_type):
  55. """ check src dataset format validity """
  56. if dataset_type == "LabelMe":
  57. anno_suffix = ".json"
  58. else:
  59. raise ConvertFailedError(
  60. message=f"数据格式转换失败!不支持{dataset_type}格式数据集。当前仅支持 LabelMe 格式。")
  61. err_msg_prefix = f"数据格式转换失败!请参考上述`{dataset_type}格式数据集示例`检查待转换数据集格式。"
  62. anno_map = {}
  63. for dst_anno, src_anno in [("instance_train.json", "train_anno_list.txt"),
  64. ("instance_val.json", "val_anno_list.txt")]:
  65. src_anno_path = os.path.join(root_dir, src_anno)
  66. if not os.path.exists(src_anno_path):
  67. if dst_anno == "instance_train.json":
  68. raise ConvertFailedError(
  69. message=f"{err_msg_prefix}保证{src_anno_path}文件存在。")
  70. continue
  71. with custom_open(src_anno_path, 'r') as f:
  72. anno_list = f.readlines()
  73. for anno_fn in anno_list:
  74. anno_fn = anno_fn.strip().split(' ')[-1]
  75. anno_path = os.path.join(root_dir, anno_fn)
  76. if not os.path.exists(anno_path):
  77. raise ConvertFailedError(
  78. message=f"{err_msg_prefix}保证\"{src_anno_path}\"中的\"{anno_fn}\"文件存在。"
  79. )
  80. anno_map[dst_anno] = src_anno_path
  81. return anno_map
  82. def convert(dataset_type, input_dir):
  83. """ convert dataset to coco format """
  84. # check format validity
  85. anno_map = check_src_dataset(input_dir, dataset_type)
  86. if dataset_type == "LabelMe":
  87. convert_labelme_dataset(input_dir, anno_map)
  88. else:
  89. raise ValueError
  90. def split_anno_list(root_dir, anno_map):
  91. """Split anno list to 80% train and 20% val """
  92. train_anno_list = []
  93. val_anno_list = []
  94. anno_list_bak = os.path.join(root_dir, "train_anno_list.txt.bak")
  95. shutil.move(anno_map["instance_train.json"], anno_list_bak),
  96. with custom_open(anno_list_bak, 'r') as f:
  97. src_anno = f.readlines()
  98. random.shuffle(src_anno)
  99. train_anno_list = src_anno[:int(len(src_anno) * 0.8)]
  100. val_anno_list = src_anno[int(len(src_anno) * 0.8):]
  101. with custom_open(os.path.join(root_dir, "train_anno_list.txt"), 'w') as f:
  102. f.writelines(train_anno_list)
  103. with custom_open(os.path.join(root_dir, "val_anno_list.txt"), 'w') as f:
  104. f.writelines(val_anno_list)
  105. anno_map["instance_train.json"] = os.path.join(root_dir,
  106. "train_anno_list.txt")
  107. anno_map["instance_val.json"] = os.path.join(root_dir, "val_anno_list.txt")
  108. msg = f"{os.path.join(root_dir,'val_anno_list.txt')}不存在,数据集已默认按照80%训练集,20%验证集划分,\
  109. 且将原始'train_anno_list.txt'重命名为'train_anno_list.txt.bak'."
  110. warning(msg)
  111. return anno_map
  112. def convert_labelme_dataset(root_dir, anno_map):
  113. """ convert dataset labeled by LabelMe to coco format """
  114. label_indexer = Indexer()
  115. img_indexer = Indexer()
  116. annotations_dir = os.path.join(root_dir, "annotations")
  117. if not os.path.exists(annotations_dir):
  118. os.makedirs(annotations_dir)
  119. # 不存在val_anno_list,对原始数据集进行划分
  120. if 'instance_val.json' not in anno_map:
  121. anno_map = split_anno_list(root_dir, anno_map)
  122. for dst_anno in anno_map:
  123. labelme2coco(label_indexer, img_indexer, root_dir, anno_map[dst_anno],
  124. os.path.join(annotations_dir, dst_anno))
  125. def labelme2coco(label_indexer, img_indexer, root_dir, anno_path, save_path):
  126. """ convert json files generated by LabelMe to coco format and save to files """
  127. import pycocotools.mask as mask_util
  128. with custom_open(anno_path, 'r') as f:
  129. json_list = f.readlines()
  130. anno_num = 0
  131. anno_list = []
  132. image_list = []
  133. info(f"Start loading json annotation files from {anno_path} ...")
  134. for json_path in tqdm(json_list):
  135. json_path = json_path.strip()
  136. assert json_path.endswith(".json"), json_path
  137. with custom_open(os.path.join(root_dir, json_path.strip()), 'r') as f:
  138. labelme_data = json.load(f)
  139. img_id = img_indexer.get_id(labelme_data['imagePath'])
  140. height = labelme_data['imageHeight']
  141. width = labelme_data['imageWidth']
  142. image_list.append({
  143. 'id': img_id,
  144. 'file_name': labelme_data['imagePath'].split('/')[-1],
  145. 'width': width,
  146. 'height': height,
  147. })
  148. for shape in labelme_data['shapes']:
  149. assert shape[
  150. 'shape_type'] == 'polygon', "Only polygon are supported."
  151. category_id = label_indexer.get_id(shape['label'])
  152. points = shape["points"]
  153. segmentation = [np.asarray(points).flatten().tolist()]
  154. mask = points_to_mask([height, width], points)
  155. mask = np.asfortranarray(mask.astype(np.uint8))
  156. mask = mask_util.encode(mask)
  157. area = float(mask_util.area(mask))
  158. bbox = mask_util.toBbox(mask).flatten().tolist()
  159. anno_num += 1
  160. anno_list.append({
  161. 'image_id': img_id,
  162. 'bbox': bbox,
  163. 'segmentation': segmentation,
  164. 'category_id': category_id,
  165. 'id': anno_num,
  166. 'iscrowd': 0,
  167. 'area': area,
  168. 'ignore': 0
  169. })
  170. category_list = label_indexer.get_list(key_name="name")
  171. data_coco = {
  172. 'images': image_list,
  173. 'categories': category_list,
  174. 'annotations': anno_list
  175. }
  176. write_json_file(data_coco, save_path)
  177. info(f"The converted annotations has been save to {save_path}.")
  178. def points_to_mask(img_shape, points):
  179. """convert polygon points to binary mask"""
  180. mask = np.zeros(img_shape[:2], dtype=np.uint8)
  181. mask = Image.fromarray(mask)
  182. draw = ImageDraw.Draw(mask)
  183. xy = [tuple(point) for point in points]
  184. assert len(xy) > 2, "Polygon must have points more than 2"
  185. draw.polygon(xy=xy, outline=1, fill=1)
  186. mask = np.asarray(mask, dtype=bool)
  187. return mask