convert_dataset.py 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157
  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 glob
  12. import json
  13. import os
  14. import os.path as osp
  15. import shutil
  16. import cv2
  17. import numpy as np
  18. from PIL import Image, ImageDraw
  19. from .....utils.file_interface import custom_open
  20. from .....utils import logging
  21. def convert_dataset(dataset_type, input_dir):
  22. """ convert to paddlex official format """
  23. if dataset_type == "LabelMe":
  24. return convert_labelme_dataset(input_dir)
  25. else:
  26. raise NotImplementedError(dataset_type)
  27. def convert_labelme_dataset(input_dir):
  28. """ convert labelme format to paddlex official format"""
  29. bg_name = "_background_"
  30. ignore_name = '__ignore__'
  31. # prepare dir
  32. output_img_dir = osp.join(input_dir, 'images')
  33. output_annot_dir = osp.join(input_dir, 'annotations')
  34. if not osp.exists(output_img_dir):
  35. os.makedirs(output_img_dir)
  36. if not osp.exists(output_annot_dir):
  37. os.makedirs(output_annot_dir)
  38. # collect class_names and set class_name_to_id
  39. class_names = []
  40. class_name_to_id = {}
  41. split_tags = ["train", "val"]
  42. for tag in split_tags:
  43. mapping_file = osp.join(input_dir, f'{tag}_anno_list.txt')
  44. with open(mapping_file, 'r') as f:
  45. label_files = [
  46. osp.join(input_dir, line.strip('\n')) for line in f.readlines()
  47. ]
  48. for label_file in label_files:
  49. with custom_open(label_file, 'r') as fp:
  50. data = json.load(fp)
  51. for shape in data['shapes']:
  52. cls_name = shape['label']
  53. if cls_name not in class_names:
  54. class_names.append(cls_name)
  55. if ignore_name in class_names:
  56. class_name_to_id[ignore_name] = 255
  57. class_names.remove(ignore_name)
  58. if bg_name in class_names:
  59. class_names.remove(bg_name)
  60. class_name_to_id[bg_name] = 0
  61. for i, name in enumerate(class_names):
  62. class_name_to_id[name] = i + 1
  63. if len(class_names) > 256:
  64. raise ValueError(
  65. f"There are {len(class_names)} categories in the annotation file, "
  66. f"exceeding 256, Not compliant with paddlex official format!")
  67. # create annotated images and copy origin images
  68. color_map = get_color_map_list(256)
  69. img_file_list = []
  70. label_file_list = []
  71. for i, label_file in enumerate(label_files):
  72. filename = osp.splitext(osp.basename(label_file))[0]
  73. annotated_img_path = osp.join(output_annot_dir, filename + '.png')
  74. with custom_open(label_file, 'r') as f:
  75. data = json.load(f)
  76. img_path = osp.join(osp.dirname(label_file), data['imagePath'])
  77. if not os.path.exists(img_path):
  78. logging.info('%s is not existed, skip this image' %
  79. img_path)
  80. continue
  81. img_name = img_path.split('/')[-1]
  82. img_file_list.append(f"images/{img_name}")
  83. label_img_name = annotated_img_path.split("/")[-1]
  84. label_file_list.append(f"annotations/{label_img_name}")
  85. img = np.asarray(cv2.imread(img_path))
  86. lbl = shape2label(
  87. img_size=img.shape,
  88. shapes=data['shapes'],
  89. class_name_mapping=class_name_to_id)
  90. lbl_pil = Image.fromarray(lbl.astype(np.uint8), mode='P')
  91. lbl_pil.putpalette(color_map)
  92. lbl_pil.save(annotated_img_path)
  93. shutil.copy(img_path, output_img_dir)
  94. with custom_open(osp.join(input_dir, f'{tag}.txt'), 'w') as fp:
  95. for img_path, lbl_path in zip(img_file_list, label_file_list):
  96. fp.write(f'{img_path} {lbl_path}\n')
  97. with custom_open(osp.join(input_dir, 'class_name.txt'), 'w') as fp:
  98. for name in class_names:
  99. fp.write(f'{name}{os.linesep}')
  100. with custom_open(osp.join(input_dir, 'class_name_to_id.txt'), 'w') as fp:
  101. for key, val in class_name_to_id.items():
  102. fp.write(f'{val}: {key}{os.linesep}')
  103. return input_dir
  104. def get_color_map_list(num_classes):
  105. """ get color map list"""
  106. num_classes += 1
  107. color_map = num_classes * [0, 0, 0]
  108. for i in range(0, num_classes):
  109. j = 0
  110. lab = i
  111. while lab:
  112. color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
  113. color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
  114. color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
  115. j += 1
  116. lab >>= 3
  117. color_map = color_map[3:]
  118. return color_map
  119. def shape2label(img_size, shapes, class_name_mapping):
  120. """ 根据输入的形状列表,将图像的标签矩阵填充为对应形状的类别编号 """
  121. label = np.zeros(img_size[:2], dtype=np.int32)
  122. for shape in shapes:
  123. points = shape['points']
  124. class_name = shape['label']
  125. label_mask = polygon2mask(img_size[:2], points)
  126. label[label_mask] = class_name_mapping[class_name]
  127. return label
  128. def polygon2mask(img_size, points):
  129. """ 将给定形状的点转换成对应的掩膜 """
  130. label_mask = Image.fromarray(np.zeros(img_size[:2], dtype=np.uint8))
  131. image_draw = ImageDraw.Draw(label_mask)
  132. points_list = [tuple(point) for point in points]
  133. assert len(points_list) > 2, ValueError(
  134. 'Polygon must have points more than 2')
  135. image_draw.polygon(xy=points_list, outline=1, fill=1)
  136. return np.array(label_mask, dtype=bool)