labelme_json_to_dataset 1.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  1. #!/usr/bin/python
  2. import argparse
  3. import json
  4. import os
  5. import os.path as osp
  6. import PIL.Image
  7. import yaml
  8. from labelme import utils
  9. def main():
  10. parser = argparse.ArgumentParser()
  11. parser.add_argument('json_file')
  12. args = parser.parse_args()
  13. json_file = args.json_file
  14. out_dir = osp.basename(json_file).replace('.', '_')
  15. out_dir = osp.join(osp.dirname(json_file), out_dir)
  16. os.mkdir(out_dir)
  17. data = json.load(open(json_file))
  18. img = utils.img_b64_to_array(data['imageData'])
  19. lbl, lbl_names = utils.labelme_shapes_to_label(img.shape, data['shapes'])
  20. lbl_viz = utils.draw_label(lbl, img, lbl_names)
  21. PIL.Image.fromarray(img).save(osp.join(out_dir, 'img.png'))
  22. PIL.Image.fromarray(lbl).save(osp.join(out_dir, 'label.png'))
  23. PIL.Image.fromarray(lbl_viz).save(osp.join(out_dir, 'label_viz.png'))
  24. info = dict(label_names=lbl_names)
  25. with open(osp.join(out_dir, 'info.yaml'), 'w') as f:
  26. yaml.safe_dump(info, f, default_flow_style=False)
  27. print('wrote data to %s' % out_dir)
  28. if __name__ == '__main__':
  29. main()