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)