data_type.py 2.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677
  1. from ipfml import processing, metrics
  2. from PIL import Image
  3. from skimage import color
  4. import random
  5. import numpy as np
  6. def get_svd_data(data_type, block):
  7. """
  8. Method which returns the data type expected
  9. """
  10. if data_type == 'lab':
  11. nb = random.randint(0, 10)
  12. block_file_path = '/tmp/' + str(nb) + '_lab_img.png'
  13. block.save(block_file_path)
  14. data = processing.get_LAB_L_SVD_s(Image.open(block_file_path))
  15. if data_type == 'mscn_revisited':
  16. nb = random.randint(0, 10)
  17. img_mscn_revisited = processing.rgb_to_mscn(block)
  18. # save tmp as img
  19. img_output = Image.fromarray(img_mscn_revisited.astype('uint8'), 'L')
  20. mscn_revisited_file_path = '/tmp/' + str(nb) + '_mscn_revisited_img.png'
  21. img_output.save(mscn_revisited_file_path)
  22. img_block = Image.open(mscn_revisited_file_path)
  23. # extract from temp image
  24. data = metrics.get_SVD_s(img_block)
  25. if data_type == 'mscn':
  26. img_gray = np.array(color.rgb2gray(np.asarray(block))*255, 'uint8')
  27. img_mscn = processing.calculate_mscn_coefficients(img_gray, 7)
  28. img_mscn_norm = processing.normalize_2D_arr(img_mscn)
  29. img_mscn_gray = np.array(img_mscn_norm*255, 'uint8')
  30. data = metrics.get_SVD_s(img_mscn_gray)
  31. if data_type == 'low_bits_6':
  32. low_bits_6 = processing.rgb_to_LAB_L_low_bits(block, 6)
  33. data = metrics.get_SVD_s(low_bits_6)
  34. if data_type == 'low_bits_5':
  35. low_bits_5 = processing.rgb_to_LAB_L_low_bits(block, 5)
  36. data = metrics.get_SVD_s(low_bits_5)
  37. if data_type == 'low_bits_4':
  38. low_bits_4 = processing.rgb_to_LAB_L_low_bits(block, 4)
  39. data = metrics.get_SVD_s(low_bits_4)
  40. if data_type == 'low_bits_3':
  41. low_bits_3 = processing.rgb_to_LAB_L_low_bits(block, 3)
  42. data = metrics.get_SVD_s(low_bits_3)
  43. if data_type == 'low_bits_2':
  44. low_bits_2 = processing.rgb_to_LAB_L_low_bits(block, 2)
  45. data = metrics.get_SVD_s(low_bits_2)
  46. if data_type == 'low_bits_4_shifted_2':
  47. data = metrics.get_SVD_s(processing.rgb_to_LAB_L_bits(block, (3, 6)))
  48. return data