1. pip install sewar
2.  
3. import tensorflow as tf
4. from tensorflow import keras
5.  
6. from tensorflow.keras.models import Sequential, Model, load_model
7. from tensorflow.keras.layers import Dense, Input, Activation
8. from tensorflow.keras.optimizers.legacy import Adam
9. from tensorflow.keras.utils import plot_model
10.  
11. from tensorflow.keras import metrics
12. from tensorflow.keras import losses
13. from skimage import metrics as mcs
14. import os
15. from sewar.full_ref import mse
16. import pandas as pd
17. from PIL import Image
18. import cv2
19. import numpy as np
20.  
21. from google.colab import drive
22. drive.mount('/content/drive')
23.  
24. df = pd.read_csv('/content/drive/My Drive/zbiór_danych/dataset_chrominance_tables.csv',index_col=[0])
25.  
26. !unzip "/content/drive/My Drive/zbiór_danych/tiny-imagenet-200-gray.zip"
27.  
28. !unzip "/content/drive/My Drive/zbiór_danych/tiny-imagenet-200-color.zip"
29.  
30. class Dataset:
31.   def __init__(self, quality):
32.     self.quality = quality
33.     self.selected_rows = df.copy()
34.     self.selected_rows.where(self.selected_rows["quality"] == quality, inplace = True)
35.     self.selected_rows = self.selected_rows[~self.selected_rows['quality'].isnull()]
36.  
37.     self.train_x = np.empty([int(self.selected_rows.shape[0]*0.9), 64*64])
38.     self.test_x = np.empty([self.selected_rows.shape[0] - int(self.selected_rows.shape[0]*0.9), 64*64])
39.     self.train_y = np.empty([int(self.selected_rows.shape[0]*0.9), 64])
40.     self.test_y = np.empty([self.selected_rows.shape[0] - int(self.selected_rows.shape[0]*0.9), 64])
41.  
42.     self.model = Sequential()
43.  
44.     self.prepare_x()
45.     self.prepare_y()
46.  
47.   def prepare_x(self):
48.     x = np.empty([self.selected_rows.shape[0],64*64])
49.     for i in range(self.selected_rows.shape[0]):
50.       imageA = cv2.imread("tiny-imagenet-200-gray/"+self.selected_rows["name"].iloc[i])
51.       grayA = cv2.cvtColor(imageA, cv2.COLOR_BGR2GRAY)
52.       x[i] = grayA.flatten()
53.  
54.     x = x.astype('float32')
55.     x /= 255
56.     self.train_x = x[:int(self.selected_rows.shape[0]*0.9)]
57.     self.test_x = x[int(self.selected_rows.shape[0]*0.9):]
58.  
59.   def prepare_y(self):
60.     y = self.selected_rows["quantization table"].to_numpy()
61.     y_array = np.empty([self.selected_rows.shape[0],64])
62.     for i in range(self.selected_rows.shape[0]):
63.       y_array[i] = [float(s) for s in y[i].strip('[]').split(',')]
64.  
65.     self.train_y = y_array[:int(self.selected_rows.shape[0]*0.9)]
66.     self.test_y = y_array[int(self.selected_rows.shape[0]*0.9):]
67.    
68.   def create_model(self):
69.     self.model.add(Input(shape=(4096,)))
70.  
71.     self.model.add(Dense(256))
72.     self.model.add(Activation('relu'))
73.  
74.     self.model.add(Dense(64))
75.     self.model.add(Activation('relu'))
76.  
77.     self.model.compile(
78.         optimizer=Adam(),
79.         loss='mse',
80.         metrics=[tf.keras.metrics.MeanSquaredError()])
81.  
82.     self.model.summary()
83.     # # visualize the model
84.     dot_img_file = '/tmp/model_1.png'
85.     plot_model(self.model, to_file=dot_img_file, show_shapes=True)
86.  
87.   def train_model(self):
88.     history = self.model.fit(self.train_x, self.train_y,
89.                     batch_size=50,
90.                     epochs=100,
91.                     verbose=1,
92.                     validation_data=(self.test_x, self.test_y))
93.    
94.   def load_weights(self, path):
95.     self.model.load_weights(path)
96.  
97.   def save_weights(self, path):
98.     self.model.save_weights(path)
99.    
100. def data_preparation(path, quality, test_color):
101.     im_pil = Image.open(path)
102.     im_pil.save("picture.jpeg", quality = quality)
103.     im_pil2 = Image.open("picture.jpeg")
104.     imageA = np.array(im_pil)
105.             # Convert RGB to BGR
106.     if(len(imageA.shape)==3):
107.         imageA = imageA[:, :, ::-1]
108.         grayA = cv2.cvtColor(imageA, cv2.COLOR_BGR2GRAY)
109.     else:
110.         grayA = imageA
111.  
112.     imageB = np.array(im_pil2)
113.             # Convert RGB to BGR
114.     if(len(imageB.shape)==3):
115.         imageB = imageB[:, :, ::-1]
116.         grayB = cv2.cvtColor(imageB, cv2.COLOR_BGR2GRAY)
117.     else:
118.         grayB = imageB
119.    
120.     (score, diff) = mcs.structural_similarity(grayA, grayB, full=True)
121.     compression = os.path.getsize("picture.jpeg")
122.     if test_color == True:
123.       score2 = mse(imageA,imageB)
124.     else:
125.       score2 = mse(grayA,grayB)
126.     print("Original compression results:")
127.     print(score)
128.     print(compression)
129.     print(score2)
130.     if test_color == True and len(im_pil2.quantization) > 1:
131.       return [score,compression,score2, im_pil2.quantization[0], im_pil2.quantization[1]]
132.     else:
133.       return [score,compression,score2, im_pil2.quantization[0], None]
134.    
135. def test(x1, path, test_color, x2 = None):
136.     qtable = list(map(int, x1))
137.     array = {0: qtable}
138.     if (x2 != None):
139.       array[1] = x2
140.     im_pil = Image.open(path)
141.     im_pil.save("picture.jpeg", qtables = array)
142.     im_pil2 = Image.open("picture.jpeg")
143.  
144.     imageA = np.array(im_pil)
145.             # Convert RGB to BGR
146.     if(len(imageA.shape)==3):
147.         imageA = imageA[:, :, ::-1]
148.         grayA = cv2.cvtColor(imageA, cv2.COLOR_BGR2GRAY)
149.     else:
150.         grayA = imageA
151.  
152.     imageB = np.array(im_pil2)
153.             # Convert RGB to BGR
154.     if(len(imageB.shape)==3):
155.         imageB = imageB[:, :, ::-1]
156.         grayB = cv2.cvtColor(imageB, cv2.COLOR_BGR2GRAY)
157.     else:
158.         grayB = imageB
159.  
160.     (score, diff) = mcs.structural_similarity(grayA, grayB, full=True)
161.     if test_color == True:
162.       score2 = mse(imageA,imageB)
163.     else:
164.       score2 = mse(grayA,grayB)                            
165.     compression = os.path.getsize("picture.jpeg")
166.     print("New compression results:")
167.     print(score)
168.     print(compression)
169.     print(score2)
170.     return [score,compression,score2]