Spaces:
Configuration error
Configuration error
File size: 3,399 Bytes
1ba539f |
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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
import numpy as np
from lib.config import cfg
from skimage.measure import compare_ssim
import os
import cv2
import imageio
class Evaluator:
def __init__(self):
self.mse = []
self.psnr = []
self.ssim = []
def psnr_metric(self, img_pred, img_gt):
mse = np.mean((img_pred - img_gt)**2)
psnr = -10 * np.log(mse) / np.log(10)
return psnr
def ssim_metric(self, rgb_pred, rgb_gt, batch):
mask_at_box = batch['mask_at_box'][0].detach().cpu().numpy()
H, W = int(cfg.H * cfg.ratio), int(cfg.W * cfg.ratio)
mask_at_box = mask_at_box.reshape(H, W)
# convert the pixels into an image
img_pred = np.zeros((H, W, 3))
img_pred[mask_at_box] = rgb_pred
img_gt = np.zeros((H, W, 3))
img_gt[mask_at_box] = rgb_gt
# crop the object region
x, y, w, h = cv2.boundingRect(mask_at_box.astype(np.uint8))
img_pred = img_pred[y:y + h, x:x + w]
img_gt = img_gt[y:y + h, x:x + w]
# compute the ssim
ssim = compare_ssim(img_pred, img_gt, multichannel=True)
return ssim
def evaluate(self, batch):
if cfg.human in [302, 313, 315]:
i = batch['i'].item() + 1
else:
i = batch['i'].item()
i = i + cfg.begin_i
cam_ind = batch['cam_ind'].item()
# obtain the image path
result_dir = 'data/result/neural_volumes/{}_nv'.format(cfg.human)
frame_dir = os.path.join(result_dir, 'frame_{}'.format(i))
gt_img_path = os.path.join(frame_dir, 'gt_{}.jpg'.format(cam_ind + 1))
pred_img_path = os.path.join(frame_dir,
'pred_{}.jpg'.format(cam_ind + 1))
mask_at_box = batch['mask_at_box'][0].detach().cpu().numpy()
H, W = int(cfg.H * cfg.ratio), int(cfg.W * cfg.ratio)
mask_at_box = mask_at_box.reshape(H, W)
# convert the pixels into an image
rgb_gt = batch['rgb'][0].detach().cpu().numpy()
img_gt = np.zeros((H, W, 3))
img_gt[mask_at_box] = rgb_gt
# gt_img_path = gt_img_path.replace('neural_volumes', 'gt')
# os.system('mkdir -p {}'.format(os.path.dirname(gt_img_path)))
# img_gt = img_gt[..., [2, 1, 0]] * 255
# cv2.imwrite(gt_img_path, img_gt)
img_pred = imageio.imread(pred_img_path).astype(np.float32) / 255.
img_pred[mask_at_box != 1] = 0
rgb_pred = img_pred[mask_at_box]
# import matplotlib.pyplot as plt
# _, (ax1, ax2) = plt.subplots(1, 2)
# ax1.imshow(img_gt)
# ax2.imshow(img_pred)
# plt.show()
# return
mse = np.mean((rgb_pred - rgb_gt)**2)
self.mse.append(mse)
psnr = self.psnr_metric(rgb_pred, rgb_gt)
self.psnr.append(psnr)
ssim = self.ssim_metric(rgb_pred, rgb_gt, batch)
self.ssim.append(ssim)
def summarize(self):
result_path = os.path.join(cfg.result_dir, 'metrics.npy')
os.system('mkdir -p {}'.format(os.path.dirname(result_path)))
metrics = {'mse': self.mse, 'psnr': self.psnr, 'ssim': self.ssim}
np.save(result_path, self.mse)
print('mse: {}'.format(np.mean(self.mse)))
print('psnr: {}'.format(np.mean(self.psnr)))
print('ssim: {}'.format(np.mean(self.ssim)))
self.mse = []
self.psnr = []
self.ssim = []
|