from numpy.random import randint
from numpy.random import rand
import numpy as np
from PIL import Image
import cv2
import os
import random
from skimage import metrics
import argparse
from sewar.full_ref import mse
import pandas as pd
import glob
import time
from google.colab import drive
drive.mount('/content/drive')
!unzip "/content/drive/My Drive/zbiór_danych/tiny-imagenet-200-gray.zip"
def fitness(x, path):
im_pil = Image.open(path)
im_pil.save("picture.jpeg", qtables = {0: list(map(int, x))} )
im_pil2 = Image.open("picture.jpeg")
imageA = np.array(im_pil)
if(len(imageA.shape)==3):
imageA = imageA[:, :, ::-1]
grayA = cv2.cvtColor(imageA, cv2.COLOR_BGR2GRAY)
else:
grayA = imageA
imageB = np.array(im_pil2)
if(len(imageB.shape)==3):
imageB = imageB[:, :, ::-1]
grayB = cv2.cvtColor(imageB, cv2.COLOR_BGR2GRAY)
else:
grayB = imageB
score = mse(grayA,grayB)
compression = os.path.getsize("picture.jpeg")/os.path.getsize(path)
return score + 30 * compression
def decode(bounds, n_bits, bitstring):
decoded = list()
largest = 2**n_bits
for i in range(len(bounds)):
start, end = i * n_bits, (i * n_bits)+n_bits
substring = bitstring[start:end]
chars = ''.join([str(s) for s in substring])
integer = int(chars, 2)
value = bounds[i][0] + int((integer/largest) * (bounds[i][1] - bounds[i][0]))
decoded.append(value)
return decoded
def selection_ranking(pop, scores, k):
pop = [y for _,y in sorted(zip(scores,pop),key=lambda x: x[0])]
pop[0:k] = pop[len(pop)-k:len(pop)]
return pop
def crossover(p1, p2, r_cross):
c1, c2 = p1.copy(), p2.copy()
if rand() < r_cross:
pt = randint(1, len(p1)-2)
c1 = p1[:pt] + p2[pt:]
c2 = p2[:pt] + p1[pt:]
return [c1, c2]
def mutation(bitstring, r_mut):
for i in range(len(bitstring)):
if rand() < r_mut:
bitstring[i] = 1 - bitstring[i]
def genetic_algorithm(fitness, bounds, n_bits, n_iter, n_pop, r_cross, r_mut, start_value, path):
new_x = [int(example * 2**n_bits / bounds[0][1]) for example in start_value]
newList = list()
for j in new_x:
newList.extend([int(x) for x in bin(j)[2:].zfill(n_bits)])
pop = [newList for _ in range(n_pop)]
best, best_eval = pop[0], fitness(decode(bounds, n_bits, pop[0]),path)
for gen in range(n_iter):
decoded = [decode(bounds, n_bits, p) for p in pop]
scores = [fitness(d,path) for d in decoded]
for i in range(n_pop):
if scores[i] < best_eval:
best, best_eval = pop[i], scores[i]
print(">%d, new best f(%s) = %f" % (gen, decoded[i], scores[i]))
selected = selection_ranking(pop, scores, 5)
children = list()
for i in range(0, n_pop, 2):
p1, p2 = selected[i], selected[i+1]
for c in crossover(p1, p2, r_cross):
mutation(c, r_mut)
children.append(c)
pop = children
return [best, best_eval]
def data_preparation(path, quality):
im_pil = Image.open(path)
im_pil.save("picture.jpeg", quality = quality)
im_pil2 = Image.open("picture.jpeg")
imageA = np.array(im_pil)
if(len(imageA.shape)==3):
imageA = imageA[:, :, ::-1]
grayA = cv2.cvtColor(imageA, cv2.COLOR_BGR2GRAY)
else:
grayA = imageA
imageB = np.array(im_pil2)
if(len(imageB.shape)==3):
imageB = imageB[:, :, ::-1]
grayB = cv2.cvtColor(imageB, cv2.COLOR_BGR2GRAY)
else:
grayB = imageB
(score_structural_similarity, diff) = metrics.structural_similarity(grayA, grayB, full=True)
compression = os.path.getsize("picture.jpeg")/os.path.getsize(path)
score_mse = mse(grayA,grayB)
return [score_structural_similarity,compression,score_mse, im_pil2.quantization[0]]
def test(x, path):
qtable = list(map(int, x))
array = {0: qtable}
im_pil = Image.open(path)
im_pil.save("picture.jpeg", qtables = array)
im_pil2 = Image.open("picture.jpeg")
imageA = np.array(im_pil)
if(len(imageA.shape)==3):
imageA = imageA[:, :, ::-1]
grayA = cv2.cvtColor(imageA, cv2.COLOR_BGR2GRAY)
else:
grayA = imageA
imageB = np.array(im_pil2)
if(len(imageB.shape)==3):
imageB = imageB[:, :, ::-1]
grayB = cv2.cvtColor(imageB, cv2.COLOR_BGR2GRAY)
else:
grayB = imageB
(score_structural_similarity, diff) = metrics.structural_similarity(grayA, grayB, full=True)
score_mse = mse(grayA,grayB)
compression = os.path.getsize("picture.jpeg")/os.path.getsize(path)
return [score_structural_similarity,compression,score_mse]
arr = os.listdir("tiny-imagenet-200-gray/")
arr.sort()
index = 0
quality_list = [70,80,85,90]
index_quality = -1
if os.path.exists('/content/drive/My Drive/zbiór_danych/quantization_tables_gray.csv'):
df = pd.read_csv('/content/drive/My Drive/zbiór_danych/quantization_tables_gray.csv',index_col=[0])
print(df)
index = arr.index(df.iloc[-1]["name"])
index_quality = quality_list.index(df.iloc[-1]["quality"])
if df.iloc[-1]["quality"] == 90:
index = index + 1
index_quality = -1
else:
df = pd.DataFrame(columns=['name','quantization table','quality', 'mse', 'compression'])
start_flag = True
print('Index: ' + str(index))
for name in arr[index:]:
if start_flag:
quality_list = quality_list[index_quality + 1:]
for quality in quality_list:
tic = time.perf_counter()
start_flag = False
print('Quality:'+str(quality))
path = "tiny-imagenet-200-gray/"+name
score_structural_similarity_original,compression_original,score_mse_original, quantization = data_preparation(path, quality)
print('Score structural similarity original: ' + str(score_structural_similarity_original))
print('Compression original: ' + str(compression_original))
print('Score mse original: ' + str(score_mse_original))
flag = False
decoded = []
while flag == False:
best, score = genetic_algorithm(fitness, bounds, n_bits, n_iter, n_pop, r_cross, r_mut, "ranking", 0.001, 1, quantization,path)
print('Done!')
decoded = decode(bounds, n_bits, best)
print('f(%s) = %f' % (decoded, score))
score_structural_similarity, compression, score_mse = test(decoded,path)
print('Score structural similarity: ' + str(score_structural_similarity))
print('Compression: ' + str(compression))
print('Score mse:' + str(score_mse))
if (score_mse < score_mse_original and compression < compression_original):
flag = True
df.loc[len(df.index)] = [name, decoded, quality, score_mse, compression]
df.to_csv('/content/drive/My Drive/zbiór_danych/quantization_tables_gray.csv')
