data.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223
  1. from ipfml import processing, metrics, utils
  2. from modules.utils.config import *
  3. from PIL import Image
  4. from skimage import color
  5. import numpy as np
  6. _scenes_names_prefix = '_scenes_names'
  7. _scenes_indices_prefix = '_scenes_indices'
  8. # store all variables from current module context
  9. context_vars = vars()
  10. def get_svd_data(data_type, block):
  11. """
  12. Method which returns the data type expected
  13. """
  14. if data_type == 'lab':
  15. block_file_path = '/tmp/lab_img.png'
  16. block.save(block_file_path)
  17. data = processing.get_LAB_L_SVD_s(Image.open(block_file_path))
  18. if data_type == 'mscn_revisited':
  19. img_mscn_revisited = processing.rgb_to_mscn(block)
  20. # save tmp as img
  21. img_output = Image.fromarray(img_mscn_revisited.astype('uint8'), 'L')
  22. mscn_revisited_file_path = '/tmp/mscn_revisited_img.png'
  23. img_output.save(mscn_revisited_file_path)
  24. img_block = Image.open(mscn_revisited_file_path)
  25. # extract from temp image
  26. data = metrics.get_SVD_s(img_block)
  27. if data_type == 'mscn':
  28. img_gray = np.array(color.rgb2gray(np.asarray(block))*255, 'uint8')
  29. img_mscn = processing.calculate_mscn_coefficients(img_gray, 7)
  30. img_mscn_norm = processing.normalize_2D_arr(img_mscn)
  31. img_mscn_gray = np.array(img_mscn_norm*255, 'uint8')
  32. data = metrics.get_SVD_s(img_mscn_gray)
  33. if data_type == 'low_bits_6':
  34. low_bits_6 = processing.rgb_to_LAB_L_low_bits(block, 6)
  35. data = metrics.get_SVD_s(low_bits_6)
  36. if data_type == 'low_bits_5':
  37. low_bits_5 = processing.rgb_to_LAB_L_low_bits(block, 5)
  38. data = metrics.get_SVD_s(low_bits_5)
  39. if data_type == 'low_bits_4':
  40. low_bits_4 = processing.rgb_to_LAB_L_low_bits(block, 4)
  41. data = metrics.get_SVD_s(low_bits_4)
  42. if data_type == 'low_bits_3':
  43. low_bits_3 = processing.rgb_to_LAB_L_low_bits(block, 3)
  44. data = metrics.get_SVD_s(low_bits_3)
  45. if data_type == 'low_bits_2':
  46. low_bits_2 = processing.rgb_to_LAB_L_low_bits(block, 2)
  47. data = metrics.get_SVD_s(low_bits_2)
  48. if data_type == 'low_bits_4_shifted_2':
  49. data = metrics.get_SVD_s(processing.rgb_to_LAB_L_bits(block, (3, 6)))
  50. if data_type == 'sub_blocks_stats':
  51. block = np.asarray(block)
  52. width, height, _= block.shape
  53. sub_width, sub_height = int(width / 4), int(height / 4)
  54. sub_blocks = processing.divide_in_blocks(block, (sub_width, sub_height))
  55. data = []
  56. for sub_b in sub_blocks:
  57. # by default use the whole lab L canal
  58. l_svd_data = np.array(processing.get_LAB_L_SVD_s(sub_b))
  59. # get information we want from svd
  60. data.append(np.mean(l_svd_data))
  61. data.append(np.median(l_svd_data))
  62. data.append(np.percentile(l_svd_data, 25))
  63. data.append(np.percentile(l_svd_data, 75))
  64. data.append(np.var(l_svd_data))
  65. area_under_curve = utils.integral_area_trapz(l_svd_data, dx=100)
  66. data.append(area_under_curve)
  67. # convert into numpy array after computing all stats
  68. data = np.asarray(data)
  69. if data_type == 'sub_blocks_stats_reduced':
  70. block = np.asarray(block)
  71. width, height, _= block.shape
  72. sub_width, sub_height = int(width / 4), int(height / 4)
  73. sub_blocks = processing.divide_in_blocks(block, (sub_width, sub_height))
  74. data = []
  75. for sub_b in sub_blocks:
  76. # by default use the whole lab L canal
  77. l_svd_data = np.array(processing.get_LAB_L_SVD_s(sub_b))
  78. # get information we want from svd
  79. data.append(np.mean(l_svd_data))
  80. data.append(np.median(l_svd_data))
  81. data.append(np.percentile(l_svd_data, 25))
  82. data.append(np.percentile(l_svd_data, 75))
  83. data.append(np.var(l_svd_data))
  84. # convert into numpy array after computing all stats
  85. data = np.asarray(data)
  86. if data_type == 'sub_blocks_area':
  87. block = np.asarray(block)
  88. width, height, _= block.shape
  89. sub_width, sub_height = int(width / 8), int(height / 8)
  90. sub_blocks = processing.divide_in_blocks(block, (sub_width, sub_height))
  91. data = []
  92. for sub_b in sub_blocks:
  93. # by default use the whole lab L canal
  94. l_svd_data = np.array(processing.get_LAB_L_SVD_s(sub_b))
  95. area_under_curve = utils.integral_area_trapz(l_svd_data, dx=50)
  96. data.append(area_under_curve)
  97. # convert into numpy array after computing all stats
  98. data = np.asarray(data)
  99. if data_type == 'sub_blocks_area_normed':
  100. block = np.asarray(block)
  101. width, height, _= block.shape
  102. sub_width, sub_height = int(width / 8), int(height / 8)
  103. sub_blocks = processing.divide_in_blocks(block, (sub_width, sub_height))
  104. data = []
  105. for sub_b in sub_blocks:
  106. # by default use the whole lab L canal
  107. l_svd_data = np.array(processing.get_LAB_L_SVD_s(sub_b))
  108. l_svd_data = utils.normalize_arr(l_svd_data)
  109. area_under_curve = utils.integral_area_trapz(l_svd_data, dx=50)
  110. data.append(area_under_curve)
  111. # convert into numpy array after computing all stats
  112. data = np.asarray(data)
  113. if data_type == 'mscn_var_4':
  114. data = _get_mscn_variance(block, (100, 100))
  115. if data_type == 'mscn_var_16':
  116. data = _get_mscn_variance(block, (50, 50))
  117. if data_type == 'mscn_var_64':
  118. data = _get_mscn_variance(block, (25, 25))
  119. return data
  120. def _get_mscn_variance(block, sub_block_size=(50, 50)):
  121. blocks = processing.divide_in_blocks(block, sub_block_size)
  122. data = []
  123. for block in blocks:
  124. mscn_coefficients = processing.get_mscn_coefficients(block)
  125. flat_coeff = mscn_coefficients.flatten()
  126. data.append(np.var(flat_coeff))
  127. return np.sort(data)
  128. def get_renderer_scenes_indices(renderer_name):
  129. if renderer_name not in renderer_choices:
  130. raise ValueError("Unknown renderer name")
  131. if renderer_name == 'all':
  132. return scenes_indices
  133. else:
  134. return context_vars[renderer_name + _scenes_indices_prefix]
  135. def get_renderer_scenes_names(renderer_name):
  136. if renderer_name not in renderer_choices:
  137. raise ValueError("Unknown renderer name")
  138. if renderer_name == 'all':
  139. return scenes_names
  140. else:
  141. return context_vars[renderer_name + _scenes_names_prefix]