Spaces:
Runtime error
Runtime error
File size: 2,120 Bytes
f670afc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
from unittest import result
from matplotlib.pyplot import hist
from torch.utils import data
from torch.utils.data.dataset import Dataset
import os,torch
from PIL import Image
import torchvision.transforms as T
from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
import torch.nn.functional as F
from imaginaire.evaluation.segmentation import get_segmentation_hist_model,get_miou,compute_hist
import lpips
from easydict import EasyDict as edict
from tqdm import tqdm
import piq
from torch.utils.data import DataLoader
from piq import FID,KID
import numpy as np
result_path = 'result/Ours-pers-sin-sty'
gt_path = 'dataset/CVACT/streetview_test'
class Dataset_img(Dataset):
def __init__(self, dir):
self.dir = dir
self.datalist = sorted(os.listdir(dir))
def __len__(self):
return len(self.datalist)
def __getitem__(self, index):
img = os.path.join(self.dir,self.datalist[index])
img = Image.open(img).convert('RGB')
img = T.ToTensor()(img)
return {'images':img}
data_gt = Dataset_img(gt_path)
data_pred = Dataset_img(result_path)
loss_fn_alex = lpips.LPIPS(net='alex',eval_mode=True).cuda()
loss_fn_squeeze = lpips.LPIPS(net='squeeze',eval_mode=True).cuda()
data_list = os.listdir(result_path)
results = edict()
results.psnr = []
results.ssim = []
results.alex = []
results.squeeze = []
results.RMSE = []
dataloader_pred = DataLoader(data_pred,batch_size=1,shuffle=False,num_workers=10)
dataloader_gt = DataLoader(data_gt,batch_size=1,shuffle=False,num_workers=10)
for i in tqdm(zip(dataloader_pred,dataloader_gt),ncols=100):
pred = i[0]['images'].cuda()
gt = i[1]['images'].cuda()
results.psnr.append(-10*F.mse_loss(pred,gt).log10().item())
results.ssim.append(ssim(pred, gt,data_range=1.).item())
results.alex.append(torch.mean(loss_fn_alex((pred*2.)-1, (2.*gt)-1)).cpu().item())
results.squeeze.append(torch.mean(loss_fn_squeeze((pred*2.)-1, (2.*gt)-1)).cpu().item())
results.RMSE.append(torch.sqrt(F.mse_loss(pred,gt)).item()*255)
for i in results:
print("%-10s"%i, ':',np.mean(results[i]))
|