data.py 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240
  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':
  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. """
  34. if data_type == 'low_bits_6':
  35. low_bits_6 = processing.rgb_to_LAB_L_low_bits(block, 6)
  36. data = metrics.get_SVD_s(low_bits_6)
  37. if data_type == 'low_bits_5':
  38. low_bits_5 = processing.rgb_to_LAB_L_low_bits(block, 5)
  39. data = metrics.get_SVD_s(low_bits_5)
  40. if data_type == 'low_bits_4':
  41. low_bits_4 = processing.rgb_to_LAB_L_low_bits(block, 4)
  42. data = metrics.get_SVD_s(low_bits_4)
  43. if data_type == 'low_bits_3':
  44. low_bits_3 = processing.rgb_to_LAB_L_low_bits(block, 3)
  45. data = metrics.get_SVD_s(low_bits_3)
  46. if data_type == 'low_bits_2':
  47. low_bits_2 = processing.rgb_to_LAB_L_low_bits(block, 2)
  48. data = metrics.get_SVD_s(low_bits_2)
  49. if data_type == 'low_bits_4_shifted_2':
  50. data = metrics.get_SVD_s(processing.rgb_to_LAB_L_bits(block, (3, 6)))
  51. if data_type == 'sub_blocks_stats':
  52. block = np.asarray(block)
  53. width, height, _= block.shape
  54. sub_width, sub_height = int(width / 4), int(height / 4)
  55. sub_blocks = processing.divide_in_blocks(block, (sub_width, sub_height))
  56. data = []
  57. for sub_b in sub_blocks:
  58. # by default use the whole lab L canal
  59. l_svd_data = np.array(processing.get_LAB_L_SVD_s(sub_b))
  60. # get information we want from svd
  61. data.append(np.mean(l_svd_data))
  62. data.append(np.median(l_svd_data))
  63. data.append(np.percentile(l_svd_data, 25))
  64. data.append(np.percentile(l_svd_data, 75))
  65. data.append(np.var(l_svd_data))
  66. area_under_curve = utils.integral_area_trapz(l_svd_data, dx=100)
  67. data.append(area_under_curve)
  68. # convert into numpy array after computing all stats
  69. data = np.asarray(data)
  70. if data_type == 'sub_blocks_stats_reduced':
  71. block = np.asarray(block)
  72. width, height, _= block.shape
  73. sub_width, sub_height = int(width / 4), int(height / 4)
  74. sub_blocks = processing.divide_in_blocks(block, (sub_width, sub_height))
  75. data = []
  76. for sub_b in sub_blocks:
  77. # by default use the whole lab L canal
  78. l_svd_data = np.array(processing.get_LAB_L_SVD_s(sub_b))
  79. # get information we want from svd
  80. data.append(np.mean(l_svd_data))
  81. data.append(np.median(l_svd_data))
  82. data.append(np.percentile(l_svd_data, 25))
  83. data.append(np.percentile(l_svd_data, 75))
  84. data.append(np.var(l_svd_data))
  85. # convert into numpy array after computing all stats
  86. data = np.asarray(data)
  87. if data_type == 'sub_blocks_area':
  88. block = np.asarray(block)
  89. width, height, _= block.shape
  90. sub_width, sub_height = int(width / 8), int(height / 8)
  91. sub_blocks = processing.divide_in_blocks(block, (sub_width, sub_height))
  92. data = []
  93. for sub_b in sub_blocks:
  94. # by default use the whole lab L canal
  95. l_svd_data = np.array(processing.get_LAB_L_SVD_s(sub_b))
  96. area_under_curve = utils.integral_area_trapz(l_svd_data, dx=50)
  97. data.append(area_under_curve)
  98. # convert into numpy array after computing all stats
  99. data = np.asarray(data)
  100. if data_type == 'sub_blocks_area_normed':
  101. block = np.asarray(block)
  102. width, height, _= block.shape
  103. sub_width, sub_height = int(width / 8), int(height / 8)
  104. sub_blocks = processing.divide_in_blocks(block, (sub_width, sub_height))
  105. data = []
  106. for sub_b in sub_blocks:
  107. # by default use the whole lab L canal
  108. l_svd_data = np.array(processing.get_LAB_L_SVD_s(sub_b))
  109. l_svd_data = utils.normalize_arr(l_svd_data)
  110. area_under_curve = utils.integral_area_trapz(l_svd_data, dx=50)
  111. data.append(area_under_curve)
  112. # convert into numpy array after computing all stats
  113. data = np.asarray(data)
  114. if data_type == 'mscn_var_4':
  115. data = _get_mscn_variance(block, (100, 100))
  116. if data_type == 'mscn_var_16':
  117. data = _get_mscn_variance(block, (50, 50))
  118. if data_type == 'mscn_var_64':
  119. data = _get_mscn_variance(block, (25, 25))
  120. if data_type == 'mscn_var_16_max':
  121. data = _get_mscn_variance(block, (50, 50))
  122. data = np.asarray(data)
  123. size = int(len(data) / 4)
  124. indices = data.argsort()[-size:][::-1]
  125. data = data[indices]
  126. if data_type == 'mscn_var_64_max':
  127. data = _get_mscn_variance(block, (25, 25))
  128. data = np.asarray(data)
  129. size = int(len(data) / 4)
  130. indices = data.argsort()[-size:][::-1]
  131. data = data[indices]
  132. return data
  133. def _get_mscn_variance(block, sub_block_size=(50, 50)):
  134. blocks = processing.divide_in_blocks(block, sub_block_size)
  135. data = []
  136. for block in blocks:
  137. mscn_coefficients = processing.get_mscn_coefficients(block)
  138. flat_coeff = mscn_coefficients.flatten()
  139. data.append(np.var(flat_coeff))
  140. return np.sort(data)
  141. def get_renderer_scenes_indices(renderer_name):
  142. if renderer_name not in renderer_choices:
  143. raise ValueError("Unknown renderer name")
  144. if renderer_name == 'all':
  145. return scenes_indices
  146. else:
  147. return context_vars[renderer_name + _scenes_indices_prefix]
  148. def get_renderer_scenes_names(renderer_name):
  149. if renderer_name not in renderer_choices:
  150. raise ValueError("Unknown renderer name")
  151. if renderer_name == 'all':
  152. return scenes_names
  153. else:
  154. return context_vars[renderer_name + _scenes_names_prefix]