pengsida/NeuralBody

1import numpy as np2from lib.config import cfg3from skimage.measure import compare_ssim4import os5import cv26import imageio7 8 9class Evaluator:10    def __init__(self):11        self.mse = []12        self.psnr = []13        self.ssim = []14 15    def psnr_metric(self, img_pred, img_gt):16        mse = np.mean((img_pred - img_gt)**2)17        psnr = -10 * np.log(mse) / np.log(10)18        return psnr19 20    def ssim_metric(self, rgb_pred, rgb_gt, batch):21        mask_at_box = batch['mask_at_box'][0].detach().cpu().numpy()22        H, W = int(cfg.H * cfg.ratio), int(cfg.W * cfg.ratio)23        mask_at_box = mask_at_box.reshape(H, W)24        # convert the pixels into an image25        img_pred = np.zeros((H, W, 3))26        img_pred[mask_at_box] = rgb_pred27        img_gt = np.zeros((H, W, 3))28        img_gt[mask_at_box] = rgb_gt29        # crop the object region30        x, y, w, h = cv2.boundingRect(mask_at_box.astype(np.uint8))31        img_pred = img_pred[y:y + h, x:x + w]32        img_gt = img_gt[y:y + h, x:x + w]33        # compute the ssim34        ssim = compare_ssim(img_pred, img_gt, multichannel=True)35        return ssim36 37    def evaluate(self, batch):38        if cfg.human in [302, 313, 315]:39            i = batch['i'].item() + 140        else:41            i = batch['i'].item()42        i = i + cfg.begin_i43        cam_ind = batch['cam_ind'].item()44 45        # obtain the image path46        result_dir = 'data/result/neural_volumes/{}_nv'.format(cfg.human)47        frame_dir = os.path.join(result_dir, 'frame_{}'.format(i))48        gt_img_path = os.path.join(frame_dir, 'gt_{}.jpg'.format(cam_ind + 1))49        pred_img_path = os.path.join(frame_dir,50                                     'pred_{}.jpg'.format(cam_ind + 1))51 52        mask_at_box = batch['mask_at_box'][0].detach().cpu().numpy()53        H, W = int(cfg.H * cfg.ratio), int(cfg.W * cfg.ratio)54        mask_at_box = mask_at_box.reshape(H, W)55 56        # convert the pixels into an image57        rgb_gt = batch['rgb'][0].detach().cpu().numpy()58        img_gt = np.zeros((H, W, 3))59        img_gt[mask_at_box] = rgb_gt60 61        # gt_img_path = gt_img_path.replace('neural_volumes', 'gt')62        # os.system('mkdir -p {}'.format(os.path.dirname(gt_img_path)))63        # img_gt = img_gt[..., [2, 1, 0]] * 25564        # cv2.imwrite(gt_img_path, img_gt)65 66        img_pred = imageio.imread(pred_img_path).astype(np.float32) / 255.67        img_pred[mask_at_box != 1] = 068        rgb_pred = img_pred[mask_at_box]69 70        # import matplotlib.pyplot as plt71        # _, (ax1, ax2) = plt.subplots(1, 2)72        # ax1.imshow(img_gt)73        # ax2.imshow(img_pred)74        # plt.show()75        # return76 77        mse = np.mean((rgb_pred - rgb_gt)**2)78        self.mse.append(mse)79 80        psnr = self.psnr_metric(rgb_pred, rgb_gt)81        self.psnr.append(psnr)82 83        ssim = self.ssim_metric(rgb_pred, rgb_gt, batch)84        self.ssim.append(ssim)85 86    def summarize(self):87        result_path = os.path.join(cfg.result_dir, 'metrics.npy')88        os.system('mkdir -p {}'.format(os.path.dirname(result_path)))89        metrics = {'mse': self.mse, 'psnr': self.psnr, 'ssim': self.ssim}90        np.save(result_path, self.mse)91        print('mse: {}'.format(np.mean(self.mse)))92        print('psnr: {}'.format(np.mean(self.psnr)))93        print('ssim: {}'.format(np.mean(self.ssim)))94        self.mse = []95        self.psnr = []96        self.ssim = []97