|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import os |
|
import pdb |
|
import shutil |
|
from collections import defaultdict |
|
|
|
import numpy as np |
|
import torch |
|
from PIL import Image |
|
from prettytable import PrettyTable |
|
from torch.utils.data import Dataset |
|
from torchmetrics.image import StructuralSimilarityIndexMeasure |
|
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity |
|
from torchvision import transforms |
|
from tqdm import tqdm |
|
from tqlt import utils as tu |
|
|
|
|
|
def write_json(path, x): |
|
"""write a json file. |
|
|
|
Args: |
|
path (str): path to write json file. |
|
x (dict): dict to write. |
|
""" |
|
import json |
|
|
|
with open(path, "w") as f: |
|
json.dump(x, f, indent=2) |
|
|
|
|
|
def img_center_padding(img_np, pad_ratio=0.2): |
|
|
|
ori_w, ori_h = img_np.shape[:2] |
|
|
|
w = round((1 + pad_ratio) * ori_w) |
|
h = round((1 + pad_ratio) * ori_h) |
|
|
|
img_pad_np = (np.ones((w, h, 3), dtype=img_np.dtype) * 255).astype(np.uint8) |
|
offset_h, offset_w = (w - img_np.shape[0]) // 2, (h - img_np.shape[1]) // 2 |
|
img_pad_np[ |
|
offset_h : offset_h + img_np.shape[0] :, offset_w : offset_w + img_np.shape[1] |
|
] = img_np |
|
|
|
return img_pad_np, offset_w, offset_h |
|
|
|
|
|
def scan_files_in_dir(directory, postfix=None, progress_bar=None) -> list: |
|
file_list = [] |
|
progress_bar = ( |
|
tqdm(total=0, desc=f"Scanning", ncols=100) |
|
if progress_bar is None |
|
else progress_bar |
|
) |
|
for entry in os.scandir(directory): |
|
if entry.is_file(): |
|
if postfix is None or os.path.splitext(entry.path)[1] in postfix: |
|
file_list.append(entry) |
|
progress_bar.total += 1 |
|
progress_bar.update(1) |
|
elif entry.is_dir(): |
|
file_list += scan_files_in_dir( |
|
entry.path, postfix=postfix, progress_bar=progress_bar |
|
) |
|
return file_list |
|
|
|
|
|
class EvalDataset(Dataset): |
|
def __init__(self, gt_folder, pred_folder, height=1024): |
|
self.gt_folder = gt_folder |
|
self.pred_folder = pred_folder |
|
self.height = height |
|
self.data = self.prepare_data() |
|
self.to_tensor = transforms.ToTensor() |
|
|
|
def extract_id_from_filename(self, filename): |
|
|
|
start_i = None |
|
for i, c in enumerate(filename): |
|
if c.isdigit(): |
|
start_i = i |
|
break |
|
if start_i is None: |
|
assert False, f"Cannot find number in filename {filename}" |
|
return filename[start_i : start_i + 8] |
|
|
|
def prepare_data(self): |
|
gt_files = scan_files_in_dir(self.gt_folder, postfix={".jpg", ".png"}) |
|
|
|
gt_dict = {self.extract_id_from_filename(file.name): file for file in gt_files} |
|
pred_files = scan_files_in_dir(self.pred_folder, postfix={".jpg", ".png"}) |
|
|
|
pred_files = list(filter(lambda x: "visualization" not in x.name, pred_files)) |
|
|
|
tuples = [] |
|
for pred_file in pred_files: |
|
pred_id = self.extract_id_from_filename(pred_file.name) |
|
if pred_id not in gt_dict: |
|
print(f"Cannot find gt file for {pred_file}") |
|
else: |
|
tuples.append((gt_dict[pred_id].path, pred_file.path)) |
|
return tuples |
|
|
|
def resize(self, img): |
|
w, h = img.size |
|
new_w = int(w * self.height / h) |
|
return img.resize((new_w, self.height), Image.LANCZOS) |
|
|
|
def __len__(self): |
|
return len(self.data) |
|
|
|
def __getitem__(self, idx): |
|
gt_path, pred_path = self.data[idx] |
|
|
|
gt, pred = self.resize(Image.open(gt_path)), self.resize(Image.open(pred_path)) |
|
if gt.height != self.height: |
|
gt = self.resize(gt) |
|
if pred.height != self.height: |
|
pred = self.resize(pred) |
|
gt = self.to_tensor(gt) |
|
pred = self.to_tensor(pred) |
|
return gt, pred |
|
|
|
|
|
def copy_resize_gt(gt_folder, height): |
|
new_folder = os.path.join( |
|
os.path.dirname(gt_folder[:-1] if gt_folder[-1] == "/" else gt_folder), |
|
f"resize_{height}", |
|
) |
|
if not os.path.exists(new_folder): |
|
os.makedirs(new_folder, exist_ok=True) |
|
for file in tqdm(os.listdir(gt_folder)): |
|
if os.path.exists(os.path.join(new_folder, file)): |
|
continue |
|
img = Image.open(os.path.join(gt_folder, file)) |
|
img = np.asarray(img) |
|
|
|
img = Image.fromarray(img) |
|
w, h = img.size |
|
img.save(os.path.join(new_folder, file)) |
|
return new_folder |
|
|
|
|
|
@torch.no_grad() |
|
def ssim(dataloader): |
|
ssim_score = 0 |
|
ssim = StructuralSimilarityIndexMeasure(data_range=1.0).to("cuda") |
|
for gt, pred in tqdm(dataloader, desc="Calculating SSIM"): |
|
batch_size = gt.size(0) |
|
gt, pred = gt.to("cuda"), pred.to("cuda") |
|
ssim_score += ssim(pred, gt) * batch_size |
|
return ssim_score / len(dataloader.dataset) |
|
|
|
|
|
@torch.no_grad() |
|
def lpips(dataloader): |
|
lpips_score = LearnedPerceptualImagePatchSimilarity(net_type="squeeze").to("cuda") |
|
score = 0 |
|
for gt, pred in tqdm(dataloader, desc="Calculating LPIPS"): |
|
batch_size = gt.size(0) |
|
pred = pred.to("cuda") |
|
gt = gt.to("cuda") |
|
|
|
gt = (gt * 2) - 1 |
|
pred = (pred * 2) - 1 |
|
score += lpips_score(gt, pred) * batch_size |
|
return score / len(dataloader.dataset) |
|
|
|
|
|
def eval(pred_folder, gt_folder): |
|
|
|
pred_sample = os.listdir(pred_folder)[0] |
|
gt_sample = os.listdir(gt_folder)[0] |
|
|
|
img = Image.open(os.path.join(pred_folder, pred_sample)) |
|
gt_img = Image.open(os.path.join(gt_folder, gt_sample)) |
|
|
|
copy_folder = None |
|
if img.height != gt_img.height: |
|
title = "--" * 30 + "Resizing GT Images to height {img.height}" + "--" * 30 |
|
print(title) |
|
gt_folder = copy_resize_gt(gt_folder, img.height) |
|
print("-" * len(title)) |
|
copy_folder = gt_folder |
|
|
|
|
|
dataset = EvalDataset(gt_folder, pred_folder, img.height) |
|
|
|
dataloader = torch.utils.data.DataLoader( |
|
dataset, |
|
batch_size=16, |
|
num_workers=0, |
|
shuffle=False, |
|
drop_last=False, |
|
) |
|
|
|
|
|
header = [] |
|
row = [] |
|
|
|
header += ["SSIM", "LPIPS"] |
|
ssim_ = ssim(dataloader).item() |
|
lpips_ = lpips(dataloader).item() |
|
row += [ssim_, lpips_] |
|
|
|
|
|
print("GT Folder : ", gt_folder) |
|
print("Pred Folder: ", pred_folder) |
|
table = PrettyTable() |
|
table.field_names = header |
|
table.add_row(row) |
|
|
|
if copy_folder is not None: |
|
shutil.rmtree(copy_folder) |
|
|
|
return ssim_, lpips_ |
|
|
|
|
|
def get_parse(): |
|
import argparse |
|
|
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("-f1", "--folder1", type=str, required=True) |
|
parser.add_argument("-f2", "--folder2", type=str, required=True) |
|
parser.add_argument("--pre", type=str, default="anigs") |
|
parser.add_argument("--pad", action="store_true") |
|
parser.add_argument("--debug", action="store_true") |
|
args = parser.parse_args() |
|
return args |
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
opt = get_parse() |
|
|
|
input_folder = opt.folder1 |
|
target_folder = opt.folder2 |
|
|
|
save_folder = os.path.join( |
|
f"./exps/metrics{opt.pre}", "psnr_results", "anigs_video" |
|
) |
|
os.makedirs(save_folder, exist_ok=True) |
|
|
|
input_folders = tu.next_folders(input_folder) |
|
|
|
results_dict = defaultdict(dict) |
|
lpips_list = [] |
|
ssim_list = [] |
|
|
|
for input_folder in input_folders: |
|
|
|
item_basename = tu.basename(input_folder) |
|
|
|
mask_item_folder = None |
|
input_item_folder = os.path.join(input_folder, "rgb") |
|
target_item_folder = os.path.join(target_folder, item_basename) |
|
|
|
if os.path.exists(input_item_folder) and os.path.exists(target_item_folder): |
|
|
|
ssim_, lpips_ = eval(input_item_folder, target_item_folder) |
|
|
|
if ssim_ == -1: |
|
continue |
|
|
|
lpips_list.append(lpips_) |
|
ssim_list.append(ssim_) |
|
|
|
results_dict[item_basename]["lpips"] = lpips_ |
|
results_dict[item_basename]["ssim"] = ssim_ |
|
if opt.debug: |
|
break |
|
print(results_dict) |
|
|
|
results_dict["all_mean"]["lpips"] = np.mean(lpips_list) |
|
results_dict["all_mean"]["ssim"] = np.mean(ssim_list) |
|
|
|
write_json(os.path.join(save_folder, "lpips_ssim.json"), results_dict) |
|
|
