data.py 8.0 KB

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