pip install sewar

import tensorflow as tf
from tensorflow import keras

from tensorflow.keras.models import Sequential, Model, load_model
from tensorflow.keras.layers import Dense, Input, Activation
from tensorflow.keras.optimizers.legacy import Adam 
from tensorflow.keras.utils import plot_model

from tensorflow.keras import metrics
from tensorflow.keras import losses
from skimage import metrics as mcs
import os
from sewar.full_ref import mse
import pandas as pd
from PIL import Image
import cv2
import numpy as np

from google.colab import drive
drive.mount('/content/drive')

df = pd.read_csv('/content/drive/My Drive/zbiór_danych/dataset_chrominance_tables.csv',index_col=[0])

!unzip "/content/drive/My Drive/zbiór_danych/tiny-imagenet-200-gray.zip"

!unzip "/content/drive/My Drive/zbiór_danych/tiny-imagenet-200-color.zip"

class Dataset:
  def __init__(self, quality):
    self.quality = quality
    self.selected_rows = df.copy()
    self.selected_rows.where(self.selected_rows["quality"] == quality, inplace = True)
    self.selected_rows = self.selected_rows[~self.selected_rows['quality'].isnull()]

    self.train_x = np.empty([int(self.selected_rows.shape[0]*0.9), 64*64])
    self.test_x = np.empty([self.selected_rows.shape[0] - int(self.selected_rows.shape[0]*0.9), 64*64])
    self.train_y = np.empty([int(self.selected_rows.shape[0]*0.9), 64])
    self.test_y = np.empty([self.selected_rows.shape[0] - int(self.selected_rows.shape[0]*0.9), 64])

    self.model = Sequential()

    self.prepare_x()
    self.prepare_y()

  def prepare_x(self):
    x = np.empty([self.selected_rows.shape[0],64*64])
    for i in range(self.selected_rows.shape[0]):
      imageA = cv2.imread("tiny-imagenet-200-gray/"+self.selected_rows["name"].iloc[i])
      grayA = cv2.cvtColor(imageA, cv2.COLOR_BGR2GRAY)
      x[i] = grayA.flatten()

    x = x.astype('float32')
    x /= 255
    self.train_x = x[:int(self.selected_rows.shape[0]*0.9)]
    self.test_x = x[int(self.selected_rows.shape[0]*0.9):]
  
  def prepare_y(self):
    y = self.selected_rows["quantization table"].to_numpy()
    y_array = np.empty([self.selected_rows.shape[0],64])
    for i in range(self.selected_rows.shape[0]):
      y_array[i] = [float(s) for s in y[i].strip('[]').split(',')]

    self.train_y = y_array[:int(self.selected_rows.shape[0]*0.9)]
    self.test_y = y_array[int(self.selected_rows.shape[0]*0.9):]
    
  def create_model(self):
    self.model.add(Input(shape=(4096,)))

    self.model.add(Dense(256))
    self.model.add(Activation('relu'))

    self.model.add(Dense(64))
    self.model.add(Activation('relu'))

    self.model.compile(
        optimizer=Adam(),
        loss='mse',
        metrics=[tf.keras.metrics.MeanSquaredError()])

    self.model.summary()
    # # visualize the model
    dot_img_file = '/tmp/model_1.png'
    plot_model(self.model, to_file=dot_img_file, show_shapes=True)

  def train_model(self):
    history = self.model.fit(self.train_x, self.train_y,
                    batch_size=50,
                    epochs=100,
                    verbose=1,
                    validation_data=(self.test_x, self.test_y))
    
  def load_weights(self, path):
    self.model.load_weights(path)
  
  def save_weights(self, path):
    self.model.save_weights(path)
    
def data_preparation(path, quality, test_color):
    im_pil = Image.open(path)
    im_pil.save("picture.jpeg", quality = quality)
    im_pil2 = Image.open("picture.jpeg")
    imageA = np.array(im_pil) 
            # Convert RGB to BGR 
    if(len(imageA.shape)==3):
        imageA = imageA[:, :, ::-1] 
        grayA = cv2.cvtColor(imageA, cv2.COLOR_BGR2GRAY)
    else:
        grayA = imageA

    imageB = np.array(im_pil2) 
            # Convert RGB to BGR 
    if(len(imageB.shape)==3):
        imageB = imageB[:, :, ::-1]
        grayB = cv2.cvtColor(imageB, cv2.COLOR_BGR2GRAY)
    else:
        grayB = imageB
    
    (score, diff) = mcs.structural_similarity(grayA, grayB, full=True)
    compression = os.path.getsize("picture.jpeg")
    if test_color == True:
      score2 = mse(imageA,imageB)
    else:
      score2 = mse(grayA,grayB)
    print("Original compression results:")
    print(score)
    print(compression)
    print(score2)
    if test_color == True and len(im_pil2.quantization) > 1:
      return [score,compression,score2, im_pil2.quantization[0], im_pil2.quantization[1]]
    else:
      return [score,compression,score2, im_pil2.quantization[0], None]
    
def test(x1, path, test_color, x2 = None):
    qtable = list(map(int, x1))
    array = {0: qtable}
    if (x2 != None):
      array[1] = x2
    im_pil = Image.open(path)
    im_pil.save("picture.jpeg", qtables = array)
    im_pil2 = Image.open("picture.jpeg")

    imageA = np.array(im_pil) 
            # Convert RGB to BGR 
    if(len(imageA.shape)==3):
        imageA = imageA[:, :, ::-1] 
        grayA = cv2.cvtColor(imageA, cv2.COLOR_BGR2GRAY)
    else:
        grayA = imageA

    imageB = np.array(im_pil2) 
            # Convert RGB to BGR 
    if(len(imageB.shape)==3):
        imageB = imageB[:, :, ::-1]
        grayB = cv2.cvtColor(imageB, cv2.COLOR_BGR2GRAY)
    else:
        grayB = imageB

    (score, diff) = mcs.structural_similarity(grayA, grayB, full=True)
    if test_color == True:
      score2 = mse(imageA,imageB)
    else:
      score2 = mse(grayA,grayB)                            
    compression = os.path.getsize("picture.jpeg")
    print("New compression results:")
    print(score)
    print(compression)
    print(score2)
    return [score,compression,score2]