data.py 2.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273
  1. import os
  2. from PIL import Image
  3. from ..config.cnn_config import *
  4. _scenes_names_prefix = '_scenes_names'
  5. _scenes_indices_prefix = '_scenes_indices'
  6. # store all variables from current module context
  7. context_vars = vars()
  8. def get_renderer_scenes_indices(renderer_name):
  9. if renderer_name not in renderer_choices:
  10. raise ValueError("Unknown renderer name")
  11. if renderer_name == 'all':
  12. return scenes_indices
  13. else:
  14. return context_vars[renderer_name + _scenes_indices_prefix]
  15. def get_renderer_scenes_names(renderer_name):
  16. if renderer_name not in renderer_choices:
  17. raise ValueError("Unknown renderer name")
  18. if renderer_name == 'all':
  19. return scenes_names
  20. else:
  21. return context_vars[renderer_name + _scenes_names_prefix]
  22. def get_scene_image_quality(img_path):
  23. # if path getting last element (image name) and extract quality
  24. img_postfix = img_path.split('/')[-1].split(scene_image_quality_separator)[-1]
  25. img_quality = img_postfix.replace(scene_image_extension, '')
  26. return int(img_quality)
  27. def get_scene_image_prefix(img_path):
  28. # if path getting last element (image name) and extract prefix
  29. img_prefix = img_path.split('/')[-1].split(scene_image_quality_separator)[0]
  30. return img_prefix
  31. def augmented_data_image(block, output_folder, prefix_image_name):
  32. rotations = [0, 90, 180, 270]
  33. img_flip_labels = ['original', 'horizontal', 'vertical', 'both']
  34. horizontal_img = block.transpose(Image.FLIP_LEFT_RIGHT)
  35. vertical_img = block.transpose(Image.FLIP_TOP_BOTTOM)
  36. both_img = block.transpose(Image.TRANSPOSE)
  37. flip_images = [block, horizontal_img, vertical_img, both_img]
  38. # rotate and flip image to increase dataset size
  39. for id, flip in enumerate(flip_images):
  40. for rotation in rotations:
  41. rotated_output_img = flip.rotate(rotation)
  42. output_reconstructed_filename = prefix_image_name + post_image_name_separator
  43. output_reconstructed_filename = output_reconstructed_filename + img_flip_labels[id] + '_' + str(rotation) + '.png'
  44. output_reconstructed_path = os.path.join(output_folder, output_reconstructed_filename)
  45. if not os.path.exists(output_reconstructed_path):
  46. rotated_output_img.save(output_reconstructed_path)