diff --git a/compression_transforms.py b/compression.py similarity index 100% rename from compression_transforms.py rename to compression.py diff --git a/evaluation.py b/evaluation.py new file mode 100644 index 0000000..49e5f9f --- /dev/null +++ b/evaluation.py @@ -0,0 +1,83 @@ +import cv2 +import os +from metrics import lpips, PSNR, ssim +import time +from tqdm import tqdm +import pandas as pd + +def imwrite(results_folder, name ,img): + cv2.imwrite(results_folder +'/' +name,img) + +def imread(data_folder, name): + img_rgb = cv2.imread(data_folder + '/' + name) + return img_rgb + +def prepare_results_folder(order, block_shape, method): + results_folder = 'results_' + f'order_{order:.3f}_'+f'blocks_{block_shape}' + + if not os.path.exists(results_folder): + os.mkdir(results_folder) + + method_folder = results_folder+'/'+method + if not os.path.exists(method_folder): + os.mkdir(method_folder) + + return method_folder + +def average_l(lst): + return sum(lst) / len(lst) + + +def test_for_parameters(compressor, data_folder, results_method_folder): + + img_list = os.listdir( data_folder ) + ssim_list = [] + psnr_list = [] + lpips_list = [] + time_encode = [] + time_decode = [] + + for imname in tqdm(img_list): + + img_rgb = imread(data_folder, imname) + start_encode = time.time() + img_rgb_code = compressor.encode_rgb(img_rgb) + end_encode = time.time() + + start_decode = time.time() + img_rgb_compress = compressor.decode_rgb(img_rgb_code) + end_decode = time.time() + + imwrite(results_method_folder, imname, img_rgb_compress) + + ssim_val = ssim(img_rgb, img_rgb_compress) + psnr_val = PSNR(img_rgb, img_rgb_compress) + lpips_val = lpips(img_rgb, img_rgb_compress) + + time_encode_val = end_encode - start_encode + time_decode_val = end_decode - start_decode + + + ssim_list.append(ssim_val) + psnr_list.append(psnr_val) + lpips_list.append(lpips_val) + time_encode.append(time_encode_val) + time_decode.append(time_decode_val) + + + ssim_avg = average_l(ssim_list) + psnr_avg = average_l(psnr_list) + lpips_avg = average_l(lpips_list) + time_encode_avg = average_l(time_encode) + time_decode_avg = average_l(time_decode) + + + df = pd.DataFrame({'image_name': (img_list+['mean']), + 'ssim': ssim_list + [ssim_avg], + 'psnr': psnr_list + [psnr_avg], + 'lpips': lpips_list + [lpips_avg], + 'time_encode':time_encode + [time_encode_avg], + 'time_decode':time_decode + [time_decode_avg]}) + + csv_path = results_method_folder +'/results_num.csv' + df.to_csv(csv_path) \ No newline at end of file diff --git a/utils.py b/metrics.py similarity index 78% rename from utils.py rename to metrics.py index 4492da2..62d7019 100644 --- a/utils.py +++ b/metrics.py @@ -1,12 +1,7 @@ -import cv2 import numpy as np -import os import scipy import torch -def average_l(lst): - return sum(lst) / len(lst) - def rgb_ssim(img0, img1, max_val, filter_size=11, filter_sigma=1.5, @@ -75,25 +70,3 @@ def rgb_lpips(np_gt, np_im, net_name, device): ssim = lambda x,y :rgb_ssim(x/255, y/255, max_val=1) lpips = lambda x,y :rgb_lpips(np.float32(x/255), np.float32(y/255), net_name='vgg', device=torch.device("cpu")) - -def imwrite(results_folder,method, name ,img): - cv2.imwrite(results_folder + '/' + method + '/' +name,img) - -def imread(data_folder, name): - img_rgb = cv2.imread(data_folder + '/' + name) - return img_rgb - -def prepare_results_folder(order, block_shape, all_methods): - results_folder = 'results_' + f'order_{order:.3f}_'+f'blocks_{block_shape}' - - if not os.path.exists(results_folder): - os.mkdir(results_folder) - - for method in all_methods: - method_folder = results_folder+'/'+method - if not os.path.exists(method_folder): - os.mkdir(method_folder) - - return results_folder - - diff --git a/test.py b/test.py index 8aadc86..acf2d46 100644 --- a/test.py +++ b/test.py @@ -1,10 +1,5 @@ -import os - -from compression_transforms import compression_factory -from utils import imwrite, imread, prepare_results_folder, PSNR, ssim, lpips, average_l -import pandas as pd -import time -from tqdm import tqdm +from compression import compression_factory +from evaluation import test_for_parameters, prepare_results_folder orders = [0.20, 0.30, 0.50] @@ -13,67 +8,12 @@ data_folder = 'datasets/kodak/' -def test_for_parameters(compressor, data_folder): - - results_folder = prepare_results_folder(compressor.order, compressor.block_size, all_methods) - - img_list = os.listdir( data_folder ) - ssim_list = [] - psnr_list = [] - lpips_list = [] - time_encode = [] - time_decode = [] - - for imname in tqdm(img_list): - - img_rgb = imread(data_folder, imname) - start_encode = time.time() - img_rgb_code = compressor.encode_rgb(img_rgb) - end_encode = time.time() - - start_decode = time.time() - img_rgb_compress = compressor.decode_rgb(img_rgb_code) - end_decode = time.time() - - imwrite(results_folder, method, imname, img_rgb_compress) - - ssim_val = ssim(img_rgb, img_rgb_compress) - psnr_val = PSNR(img_rgb, img_rgb_compress) - lpips_val = lpips(img_rgb, img_rgb_compress) - - time_encode_val = end_encode - start_encode - time_decode_val = end_decode - start_decode - - - ssim_list.append(ssim_val) - psnr_list.append(psnr_val) - lpips_list.append(lpips_val) - time_encode.append(time_encode_val) - time_decode.append(time_decode_val) - - - ssim_avg = average_l(ssim_list) - psnr_avg = average_l(psnr_list) - lpips_avg = average_l(lpips_list) - time_encode_avg = average_l(time_encode) - time_decode_avg = average_l(time_decode) - - - df = pd.DataFrame({'image_name': (img_list+['mean']), - 'ssim': ssim_list + [ssim_avg], - 'psnr': psnr_list + [psnr_avg], - 'lpips': lpips_list + [lpips_avg], - 'time_encode':time_encode + [time_encode_avg], - 'time_decode':time_decode + [time_decode_avg]}) - - csv_path = results_folder + '/' + method +'/results_num.csv' - df.to_csv(csv_path) - for block in block_shapes: for order in orders: for method in all_methods: + results_folder = prepare_results_folder(order, block, method) compressor = compression_factory(method, order, block) print(f'Order: {order}, method: {method}, block: {block}') - test_for_parameters(compressor, data_folder) + test_for_parameters(compressor, data_folder, results_folder)