data.py 7.5 KB

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