ego_to_json_4.py 1.3 KB

12345678910111213141516171819202122232425262728293031
  1. def create_annotations(directory, img_dir):
  2. annotations = {}
  3. for dir_name in os.listdir(directory):
  4. if not (dir_name == "images"):
  5. for _, _, files in os.walk(os.path.join(directory, dir_name)):
  6. mat = scipy.io.loadmat(os.path.join(directory, dir_name, ANNOTATION_FILE))
  7. for i, img_file in enumerate(sorted(files)):
  8. if not (img_file.endswith(".mat")):
  9. new_img_file = dir_name + "_" + img_file
  10. image = {
  11. "name": new_img_file,
  12. "objects": []
  13. }
  14. for segmentation in mat["polygons"][0][i]:
  15. if segmentation.any():
  16. image["objects"].append(segmentation.tolist())
  17. annotations[new_img_file] = image
  18. shutil.move(os.path.join(directory, dir_name, img_file), os.path.join(img_dir, new_img_file))
  19. with open(os.path.join(directory, SAVE_FILE), 'w') as output_json_file:
  20. json.dump(annotations, output_json_file)
  21. for dir_name in os.listdir(directory):
  22. if not (dir_name == "images" or dir_name == "annotations.json"):
  23. shutil.rmtree(os.path.join(directory, dir_name))