Spaces:
Configuration error
Configuration error
| 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 = [] | |