x2coco.py 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371
  1. #!/usr/bin/env python
  2. # coding: utf-8
  3. # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  4. #
  5. # Licensed under the Apache License, Version 2.0 (the "License");
  6. # you may not use this file except in compliance with the License.
  7. # You may obtain a copy of the License at
  8. #
  9. # http://www.apache.org/licenses/LICENSE-2.0
  10. #
  11. # Unless required by applicable law or agreed to in writing, software
  12. # distributed under the License is distributed on an "AS IS" BASIS,
  13. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14. # See the License for the specific language governing permissions and
  15. # limitations under the License.
  16. import cv2
  17. import json
  18. import os
  19. import os.path as osp
  20. import platform
  21. import shutil
  22. import numpy as np
  23. import PIL.ImageDraw
  24. from .base import MyEncoder, is_pic, get_encoding
  25. from paddlex.utils import path_normalization
  26. class X2COCO(object):
  27. def __init__(self):
  28. self.images_list = []
  29. self.categories_list = []
  30. self.annotations_list = []
  31. def generate_categories_field(self, label, labels_list):
  32. category = {}
  33. category["supercategory"] = "component"
  34. category["id"] = len(labels_list) + 1
  35. category["name"] = label
  36. return category
  37. def generate_rectangle_anns_field(self, points, label, image_id, object_id, label_to_num):
  38. annotation = {}
  39. seg_points = np.asarray(points).copy()
  40. seg_points[1, :] = np.asarray(points)[2, :]
  41. seg_points[2, :] = np.asarray(points)[1, :]
  42. annotation["segmentation"] = [list(seg_points.flatten())]
  43. annotation["iscrowd"] = 0
  44. annotation["image_id"] = image_id + 1
  45. annotation["bbox"] = list(
  46. map(float, [
  47. points[0][0], points[0][1], points[1][0] - points[0][0], points[1][
  48. 1] - points[0][1]
  49. ]))
  50. annotation["area"] = annotation["bbox"][2] * annotation["bbox"][3]
  51. annotation["category_id"] = label_to_num[label]
  52. annotation["id"] = object_id + 1
  53. return annotation
  54. def convert(self, image_dir, json_dir, dataset_save_dir):
  55. """转换。
  56. Args:
  57. image_dir (str): 图像文件存放的路径。
  58. json_dir (str): 与每张图像对应的json文件的存放路径。
  59. dataset_save_dir (str): 转换后数据集存放路径。
  60. """
  61. assert osp.exists(image_dir), "he image folder does not exist!"
  62. assert osp.exists(json_dir), "The json folder does not exist!"
  63. assert osp.exists(dataset_save_dir), "The save folder does not exist!"
  64. # Convert the image files.
  65. new_image_dir = osp.join(dataset_save_dir, "JPEGImages")
  66. if osp.exists(new_image_dir):
  67. shutil.rmtree(new_image_dir)
  68. os.makedirs(new_image_dir)
  69. for img_name in os.listdir(image_dir):
  70. if is_pic(img_name):
  71. shutil.copyfile(
  72. osp.join(image_dir, img_name),
  73. osp.join(new_image_dir, img_name))
  74. # Convert the json files.
  75. self.parse_json(new_image_dir, json_dir)
  76. coco_data = {}
  77. coco_data["images"] = self.images_list
  78. coco_data["categories"] = self.categories_list
  79. coco_data["annotations"] = self.annotations_list
  80. json_path = osp.join(dataset_save_dir, "annotations.json")
  81. json.dump(
  82. coco_data,
  83. open(json_path, "w"),
  84. indent=4,
  85. cls=MyEncoder)
  86. class LabelMe2COCO(X2COCO):
  87. """将使用LabelMe标注的数据集转换为COCO数据集。
  88. """
  89. def __init__(self):
  90. super(LabelMe2COCO, self).__init__()
  91. def generate_images_field(self, json_info, image_id):
  92. image = {}
  93. image["height"] = json_info["imageHeight"]
  94. image["width"] = json_info["imageWidth"]
  95. image["id"] = image_id + 1
  96. json_info["imagePath"] = path_normalization(json_info["imagePath"])
  97. image["file_name"] = osp.split(json_info["imagePath"])[-1]
  98. return image
  99. def generate_polygon_anns_field(self, height, width,
  100. points, label, image_id,
  101. object_id, label_to_num):
  102. annotation = {}
  103. annotation["segmentation"] = [list(np.asarray(points).flatten())]
  104. annotation["iscrowd"] = 0
  105. annotation["image_id"] = image_id + 1
  106. annotation["bbox"] = list(map(float, self.get_bbox(height, width, points)))
  107. annotation["area"] = annotation["bbox"][2] * annotation["bbox"][3]
  108. annotation["category_id"] = label_to_num[label]
  109. annotation["id"] = object_id + 1
  110. return annotation
  111. def get_bbox(self, height, width, points):
  112. polygons = points
  113. mask = np.zeros([height, width], dtype=np.uint8)
  114. mask = PIL.Image.fromarray(mask)
  115. xy = list(map(tuple, polygons))
  116. PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
  117. mask = np.array(mask, dtype=bool)
  118. index = np.argwhere(mask == 1)
  119. rows = index[:, 0]
  120. clos = index[:, 1]
  121. left_top_r = np.min(rows)
  122. left_top_c = np.min(clos)
  123. right_bottom_r = np.max(rows)
  124. right_bottom_c = np.max(clos)
  125. return [
  126. left_top_c, left_top_r, right_bottom_c - left_top_c,
  127. right_bottom_r - left_top_r
  128. ]
  129. def parse_json(self, img_dir, json_dir):
  130. image_id = -1
  131. object_id = -1
  132. labels_list = []
  133. label_to_num = {}
  134. for img_file in os.listdir(img_dir):
  135. img_name_part = osp.splitext(img_file)[0]
  136. json_file = osp.join(json_dir, img_name_part + ".json")
  137. if not osp.exists(json_file):
  138. os.remove(osp.join(image_dir, img_file))
  139. continue
  140. image_id = image_id + 1
  141. with open(json_file, mode='r', \
  142. encoding=get_encoding(json_file)) as j:
  143. json_info = json.load(j)
  144. img_info = self.generate_images_field(json_info, image_id)
  145. self.images_list.append(img_info)
  146. for shapes in json_info["shapes"]:
  147. object_id = object_id + 1
  148. label = shapes["label"]
  149. if label not in labels_list:
  150. self.categories_list.append(\
  151. self.generate_categories_field(label, labels_list))
  152. labels_list.append(label)
  153. label_to_num[label] = len(labels_list)
  154. points = shapes["points"]
  155. p_type = shapes["shape_type"]
  156. if p_type == "polygon":
  157. self.annotations_list.append(
  158. self.generate_polygon_anns_field(json_info["imageHeight"], json_info[
  159. "imageWidth"], points, label, image_id,
  160. object_id, label_to_num))
  161. if p_type == "rectangle":
  162. points.append([points[0][0], points[1][1]])
  163. points.append([points[1][0], points[0][1]])
  164. self.annotations_list.append(
  165. self.generate_rectangle_anns_field(points, label, image_id,
  166. object_id, label_to_num))
  167. class EasyData2COCO(X2COCO):
  168. """将使用EasyData标注的检测或分割数据集转换为COCO数据集。
  169. """
  170. def __init__(self):
  171. super(EasyData2COCO, self).__init__()
  172. def generate_images_field(self, img_path, image_id):
  173. image = {}
  174. img = cv2.imread(img_path)
  175. image["height"] = img.shape[0]
  176. image["width"] = img.shape[1]
  177. image["id"] = image_id + 1
  178. img_path = path_normalization(img_path)
  179. image["file_name"] = osp.split(img_path)[-1]
  180. return image
  181. def generate_polygon_anns_field(self, points, segmentation,
  182. label, image_id, object_id,
  183. label_to_num):
  184. annotation = {}
  185. annotation["segmentation"] = segmentation
  186. annotation["iscrowd"] = 1 if len(segmentation) > 1 else 0
  187. annotation["image_id"] = image_id + 1
  188. annotation["bbox"] = list(map(float, [
  189. points[0][0], points[0][1], points[1][0] - points[0][0], points[1][
  190. 1] - points[0][1]
  191. ]))
  192. annotation["area"] = annotation["bbox"][2] * annotation["bbox"][3]
  193. annotation["category_id"] = label_to_num[label]
  194. annotation["id"] = object_id + 1
  195. return annotation
  196. def parse_json(self, img_dir, json_dir):
  197. from pycocotools.mask import decode
  198. image_id = -1
  199. object_id = -1
  200. labels_list = []
  201. label_to_num = {}
  202. for img_file in os.listdir(img_dir):
  203. img_name_part = osp.splitext(img_file)[0]
  204. json_file = osp.join(json_dir, img_name_part + ".json")
  205. if not osp.exists(json_file):
  206. os.remove(osp.join(image_dir, img_file))
  207. continue
  208. image_id = image_id + 1
  209. with open(json_file, mode='r', \
  210. encoding=get_encoding(json_file)) as j:
  211. json_info = json.load(j)
  212. img_info = self.generate_images_field(osp.join(img_dir, img_file), image_id)
  213. self.images_list.append(img_info)
  214. for shapes in json_info["labels"]:
  215. object_id = object_id + 1
  216. label = shapes["name"]
  217. if label not in labels_list:
  218. self.categories_list.append(\
  219. self.generate_categories_field(label, labels_list))
  220. labels_list.append(label)
  221. label_to_num[label] = len(labels_list)
  222. points = [[shapes["x1"], shapes["y1"]],
  223. [shapes["x2"], shapes["y2"]]]
  224. if "mask" not in shapes:
  225. points.append([points[0][0], points[1][1]])
  226. points.append([points[1][0], points[0][1]])
  227. self.annotations_list.append(
  228. self.generate_rectangle_anns_field(points, label, image_id,
  229. object_id, label_to_num))
  230. else:
  231. mask_dict = {}
  232. mask_dict['size'] = [img_info["height"], img_info["width"]]
  233. mask_dict['counts'] = shapes['mask'].encode()
  234. mask = decode(mask_dict)
  235. contours, hierarchy = cv2.findContours(
  236. (mask).astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
  237. segmentation = []
  238. for contour in contours:
  239. contour_list = contour.flatten().tolist()
  240. if len(contour_list) > 4:
  241. segmentation.append(contour_list)
  242. self.annotations_list.append(
  243. self.generate_polygon_anns_field(points, segmentation, label, image_id, object_id,
  244. label_to_num))
  245. class JingLing2COCO(X2COCO):
  246. """将使用EasyData标注的检测或分割数据集转换为COCO数据集。
  247. """
  248. def __init__(self):
  249. super(JingLing2COCO, self).__init__()
  250. def generate_images_field(self, json_info, image_id):
  251. image = {}
  252. image["height"] = json_info["size"]["height"]
  253. image["width"] = json_info["size"]["width"]
  254. image["id"] = image_id + 1
  255. win_sep = "\\"
  256. other_sep = "/"
  257. if platform.system() == "Windows":
  258. json_info["path"] = win_sep.join(json_info["path"].split(other_sep))
  259. else:
  260. json_info["path"] = other_sep.join(json_info["path"].split(win_sep))
  261. image["file_name"] = osp.split(json_info["path"])[-1]
  262. return image
  263. def generate_polygon_anns_field(self, height, width,
  264. points, label, image_id,
  265. object_id, label_to_num):
  266. annotation = {}
  267. annotation["segmentation"] = [list(np.asarray(points).flatten())]
  268. annotation["iscrowd"] = 0
  269. annotation["image_id"] = image_id + 1
  270. annotation["bbox"] = list(map(float, self.get_bbox(height, width, points)))
  271. annotation["area"] = annotation["bbox"][2] * annotation["bbox"][3]
  272. annotation["category_id"] = label_to_num[label]
  273. annotation["id"] = object_id + 1
  274. return annotation
  275. def get_bbox(self, height, width, points):
  276. polygons = points
  277. mask = np.zeros([height, width], dtype=np.uint8)
  278. mask = PIL.Image.fromarray(mask)
  279. xy = list(map(tuple, polygons))
  280. PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
  281. mask = np.array(mask, dtype=bool)
  282. index = np.argwhere(mask == 1)
  283. rows = index[:, 0]
  284. clos = index[:, 1]
  285. left_top_r = np.min(rows)
  286. left_top_c = np.min(clos)
  287. right_bottom_r = np.max(rows)
  288. right_bottom_c = np.max(clos)
  289. return [
  290. left_top_c, left_top_r, right_bottom_c - left_top_c,
  291. right_bottom_r - left_top_r
  292. ]
  293. def parse_json(self, img_dir, json_dir):
  294. image_id = -1
  295. object_id = -1
  296. labels_list = []
  297. label_to_num = {}
  298. for img_file in os.listdir(img_dir):
  299. img_name_part = osp.splitext(img_file)[0]
  300. json_file = osp.join(json_dir, img_name_part + ".json")
  301. if not osp.exists(json_file):
  302. os.remove(osp.join(image_dir, img_file))
  303. continue
  304. image_id = image_id + 1
  305. with open(json_file, mode='r', \
  306. encoding=get_encoding(json_file)) as j:
  307. json_info = json.load(j)
  308. img_info = self.generate_images_field(json_info, image_id)
  309. self.images_list.append(img_info)
  310. anns_type = "bndbox"
  311. for i, obj in enumerate(json_info["outputs"]["object"]):
  312. if i == 0:
  313. if "polygon" in obj:
  314. anns_type = "polygon"
  315. else:
  316. if anns_type not in obj:
  317. continue
  318. object_id = object_id + 1
  319. label = obj["name"]
  320. if label not in labels_list:
  321. self.categories_list.append(\
  322. self.generate_categories_field(label, labels_list))
  323. labels_list.append(label)
  324. label_to_num[label] = len(labels_list)
  325. if anns_type == "polygon":
  326. points = []
  327. for j in range(int(len(obj["polygon"]) / 2.0)):
  328. points.append([obj["polygon"]["x" + str(j + 1)],
  329. obj["polygon"]["y" + str(j + 1)]])
  330. self.annotations_list.append(
  331. self.generate_polygon_anns_field(json_info["size"]["height"],
  332. json_info["size"]["width"],
  333. points,
  334. label,
  335. image_id,
  336. object_id,
  337. label_to_num))
  338. if anns_type == "bndbox":
  339. points = []
  340. points.append([obj["bndbox"]["xmin"], obj["bndbox"]["ymin"]])
  341. points.append([obj["bndbox"]["xmax"], obj["bndbox"]["ymax"]])
  342. points.append([obj["bndbox"]["xmin"], obj["bndbox"]["ymax"]])
  343. points.append([obj["bndbox"]["xmax"], obj["bndbox"]["ymin"]])
  344. self.annotations_list.append(
  345. self.generate_rectangle_anns_field(points, label, image_id,
  346. object_id, label_to_num))