labelme2voc.py 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109
  1. #!/usr/bin/env python
  2. from __future__ import print_function
  3. import argparse
  4. import glob
  5. import json
  6. import os
  7. import os.path as osp
  8. import numpy as np
  9. import PIL.Image
  10. import labelme
  11. def main():
  12. parser = argparse.ArgumentParser(
  13. formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  14. parser.add_argument('labels_file')
  15. parser.add_argument('in_dir')
  16. parser.add_argument('out_dir')
  17. args = parser.parse_args()
  18. if osp.exists(args.out_dir):
  19. print('Output directory already exists:', args.out_dir)
  20. quit(1)
  21. os.makedirs(args.out_dir)
  22. os.makedirs(osp.join(args.out_dir, 'JPEGImages'))
  23. os.makedirs(osp.join(args.out_dir, 'SegmentationClass'))
  24. os.makedirs(osp.join(args.out_dir, 'SegmentationClassPNG'))
  25. os.makedirs(osp.join(args.out_dir, 'SegmentationClassVisualization'))
  26. os.makedirs(osp.join(args.out_dir, 'SegmentationObject'))
  27. os.makedirs(osp.join(args.out_dir, 'SegmentationObjectPNG'))
  28. os.makedirs(osp.join(args.out_dir, 'SegmentationObjectVisualization'))
  29. print('Creating dataset:', args.out_dir)
  30. class_names = []
  31. class_name_to_id = {}
  32. for i, line in enumerate(open(args.labels_file).readlines()):
  33. class_id = i - 1 # starts with -1
  34. class_name = line.strip()
  35. class_name_to_id[class_name] = class_id
  36. if class_id == -1:
  37. assert class_name == '__ignore__'
  38. continue
  39. elif class_id == 0:
  40. assert class_name == '_background_'
  41. class_names.append(class_name)
  42. class_names = tuple(class_names)
  43. print('class_names:', class_names)
  44. out_class_names_file = osp.join(args.out_dir, 'class_names.txt')
  45. with open(out_class_names_file, 'w') as f:
  46. f.writelines('\n'.join(class_names))
  47. print('Saved class_names:', out_class_names_file)
  48. colormap = labelme.utils.label_colormap(255)
  49. for label_file in glob.glob(osp.join(args.in_dir, '*.json')):
  50. print('Generating dataset from:', label_file)
  51. with open(label_file) as f:
  52. base = osp.splitext(osp.basename(label_file))[0]
  53. out_img_file = osp.join(
  54. args.out_dir, 'JPEGImages', base + '.jpg')
  55. out_cls_file = osp.join(
  56. args.out_dir, 'SegmentationClass', base + '.npy')
  57. out_clsp_file = osp.join(
  58. args.out_dir, 'SegmentationClassPNG', base + '.png')
  59. out_clsv_file = osp.join(
  60. args.out_dir, 'SegmentationClassVisualization', base + '.jpg')
  61. out_ins_file = osp.join(
  62. args.out_dir, 'SegmentationObject', base + '.npy')
  63. out_insp_file = osp.join(
  64. args.out_dir, 'SegmentationObjectPNG', base + '.png')
  65. out_insv_file = osp.join(
  66. args.out_dir, 'SegmentationObjectVisualization', base + '.jpg')
  67. data = json.load(f)
  68. img_file = osp.join(osp.dirname(label_file), data['imagePath'])
  69. img = np.asarray(PIL.Image.open(img_file))
  70. PIL.Image.fromarray(img).save(out_img_file)
  71. cls, ins = labelme.utils.shapes_to_label(
  72. img_shape=img.shape,
  73. shapes=data['shapes'],
  74. label_name_to_value=class_name_to_id,
  75. type='instance',
  76. )
  77. ins[cls == -1] = 0 # ignore it.
  78. # class label
  79. labelme.utils.lblsave(out_clsp_file, cls)
  80. np.save(out_cls_file, cls)
  81. clsv = labelme.utils.draw_label(
  82. cls, img, class_names, colormap=colormap)
  83. PIL.Image.fromarray(clsv).save(out_clsv_file)
  84. # instance label
  85. labelme.utils.lblsave(out_insp_file, ins)
  86. np.save(out_ins_file, ins)
  87. instance_ids = np.unique(ins)
  88. instance_names = [str(i) for i in range(max(instance_ids) + 1)]
  89. insv = labelme.utils.draw_label(ins, img, instance_names)
  90. PIL.Image.fromarray(insv).save(out_insv_file)
  91. if __name__ == '__main__':
  92. main()