data.py 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. # main imports
  2. import os
  3. import numpy as np
  4. import random
  5. # image processing imports
  6. from PIL import Image
  7. # modules imports
  8. from ..config.cnn_config import *
  9. _scenes_names_prefix = '_scenes_names'
  10. _scenes_indices_prefix = '_scenes_indices'
  11. # store all variables from current module context
  12. context_vars = vars()
  13. def get_renderer_scenes_indices(renderer_name):
  14. if renderer_name not in renderer_choices:
  15. raise ValueError("Unknown renderer name")
  16. if renderer_name == 'all':
  17. return scenes_indices
  18. else:
  19. return context_vars[renderer_name + _scenes_indices_prefix]
  20. def get_renderer_scenes_names(renderer_name):
  21. if renderer_name not in renderer_choices:
  22. raise ValueError("Unknown renderer name")
  23. if renderer_name == 'all':
  24. return scenes_names
  25. else:
  26. return context_vars[renderer_name + _scenes_names_prefix]
  27. def get_scene_image_quality(img_path):
  28. # if path getting last element (image name) and extract quality
  29. img_postfix = img_path.split('/')[-1].split(scene_image_quality_separator)[-1]
  30. img_quality = img_postfix.replace(scene_image_extension, '')
  31. return int(img_quality)
  32. def get_scene_image_postfix(img_path):
  33. # if path getting last element (image name) and extract quality
  34. img_postfix = img_path.split('/')[-1].split(scene_image_quality_separator)[-1]
  35. img_quality = img_postfix.replace(scene_image_extension, '')
  36. return img_quality
  37. def get_scene_image_prefix(img_path):
  38. # if path getting last element (image name) and extract prefix
  39. img_prefix = img_path.split('/')[-1].split(scene_image_quality_separator)[0]
  40. return img_prefix