def _test(metric_device): y_pred = torch.rand(offset * idist.get_world_size(), 3, 28, 28, dtype=torch.float, device=device) y = y_pred * 0.65 def update(engine, i): return ( y_pred[i * s + offset * rank:(i + 1) * s + offset * rank], y[i * s + offset * rank:(i + 1) * s + offset * rank], ) engine = Engine(update) SSIM(data_range=1.0, device=metric_device).attach(engine, "ssim") data = list(range(n_iters)) engine.run(data=data, max_epochs=1) assert "ssim" in engine.state.metrics res = engine.state.metrics["ssim"] np_pred = y_pred.permute(0, 2, 3, 1).cpu().numpy() np_true = np_pred * 0.65 true_res = ski_ssim(np_pred, np_true, win_size=11, multichannel=True, gaussian_weights=True, data_range=1.0) assert pytest.approx(res, abs=tol) == true_res engine = Engine(update) SSIM(data_range=1.0, gaussian=False, kernel_size=7, device=metric_device).attach(engine, "ssim") data = list(range(n_iters)) engine.run(data=data, max_epochs=1) assert "ssim" in engine.state.metrics res = engine.state.metrics["ssim"] np_pred = y_pred.permute(0, 2, 3, 1).cpu().numpy() np_true = np_pred * 0.65 true_res = ski_ssim(np_pred, np_true, win_size=7, multichannel=True, gaussian_weights=False, data_range=1.0) assert pytest.approx(res, abs=tol) == true_res
def test_ssim(device, shape, kernel_size, gaussian, use_sample_covariance): y_pred = torch.rand(shape, device=device) y = y_pred * 0.8 sigma = 1.5 data_range = 1.0 ssim = SSIM(data_range=data_range, sigma=sigma, device=device) ssim.update((y_pred, y)) ignite_ssim = ssim.compute() skimg_pred = y_pred.cpu().numpy() skimg_y = skimg_pred * 0.8 skimg_ssim = ski_ssim( skimg_pred, skimg_y, win_size=kernel_size, sigma=sigma, channel_axis=1, gaussian_weights=gaussian, data_range=data_range, use_sample_covariance=use_sample_covariance, ) assert isinstance(ignite_ssim, float) assert np.allclose(ignite_ssim, skimg_ssim, atol=7e-5)
def test_ssim(): device = "cuda" if torch.cuda.is_available() else "cpu" ssim = SSIM(data_range=1.0, device=device) y_pred = torch.rand(16, 3, 64, 64, device=device) y = y_pred * 0.65 ssim.update((y_pred, y)) np_pred = y_pred.permute(0, 2, 3, 1).cpu().numpy() np_y = np_pred * 0.65 np_ssim = ski_ssim(np_pred, np_y, win_size=11, multichannel=True, gaussian_weights=True, data_range=1.0) assert isinstance(ssim.compute(), torch.Tensor) assert torch.allclose(ssim.compute(), torch.tensor(np_ssim, dtype=torch.float64, device=device), atol=1e-4) device = "cuda" if torch.cuda.is_available() else "cpu" ssim = SSIM(data_range=1.0, gaussian=False, kernel_size=7, device=device) y_pred = torch.rand(16, 3, 227, 227, device=device) y = y_pred * 0.65 ssim.update((y_pred, y)) np_pred = y_pred.permute(0, 2, 3, 1).cpu().numpy() np_y = np_pred * 0.65 np_ssim = ski_ssim(np_pred, np_y, win_size=7, multichannel=True, gaussian_weights=False, data_range=1.0) assert isinstance(ssim.compute(), torch.Tensor) assert torch.allclose(ssim.compute(), torch.tensor(np_ssim, dtype=torch.float64, device=device), atol=1e-4)
def test_ssim(size, channel, plus, multichannel): device = "cuda" if torch.cuda.is_available() else "cpu" pred = torch.rand(1, channel, size, size, device=device) target = pred + plus ssim_idx = ssim(pred, target) np_pred = np.random.rand(size, size, channel) if multichannel is False: np_pred = np_pred[:, :, 0] np_target = np.add(np_pred, plus) sk_ssim_idx = ski_ssim(np_pred, np_target, win_size=11, multichannel=multichannel, gaussian_weights=True) assert torch.allclose(ssim_idx, torch.tensor(sk_ssim_idx, dtype=torch.float, device=device), atol=1e-2, rtol=1e-2) ssim_idx = ssim(pred, pred) assert torch.allclose(ssim_idx, torch.tensor(1.0, device=device))
def test_ssim(size, channel, coef, multichannel): device = "cuda" if torch.cuda.is_available() else "cpu" pred = torch.rand(size, channel, size, size, device=device) target = pred * coef ssim_idx = ssim(pred, target, data_range=1.0) np_pred = pred.permute(0, 2, 3, 1).cpu().numpy() if multichannel is False: np_pred = np_pred[:, :, :, 0] np_target = np.multiply(np_pred, coef) sk_ssim_idx = ski_ssim( np_pred, np_target, win_size=11, multichannel=multichannel, gaussian_weights=True, data_range=1.0 ) assert torch.allclose(ssim_idx, torch.tensor(sk_ssim_idx, dtype=torch.float, device=device), atol=1e-4) ssim_idx = ssim(pred, pred) assert torch.allclose(ssim_idx, torch.tensor(1.0, device=device))
def _test_ssim(y_pred, y, data_range, kernel_size, sigma, gaussian, use_sample_covariance, device): atol = 7e-5 ssim = SSIM(data_range=data_range, sigma=sigma, device=device) ssim.update((y_pred, y)) ignite_ssim = ssim.compute() skimg_pred = y_pred.cpu().numpy() skimg_y = skimg_pred * 0.8 skimg_ssim = ski_ssim( skimg_pred, skimg_y, win_size=kernel_size, sigma=sigma, channel_axis=1, gaussian_weights=gaussian, data_range=data_range, use_sample_covariance=use_sample_covariance, ) assert isinstance(ignite_ssim, torch.Tensor) assert ignite_ssim.dtype == torch.float64 assert ignite_ssim.device == torch.device(device) assert np.allclose(ignite_ssim.cpu().numpy(), skimg_ssim, atol=atol)
def test(model, valid_loader, configs, epoch, folder, is_sunspots): print('====> ', 'test for epoch %s' % str(epoch), datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')) res_path = os.path.join(folder, str(epoch)) if not os.path.exists(res_path): os.mkdir(res_path) avg_mse = 0 img_mse, ssim = [], [] batch_num = len(valid_loader) test_size = batch_num * configs.batch_size for i in range(configs.predict_length): img_mse.append(0) ssim.append(0) with torch.no_grad(): real_input_flag = generator_input_flag(configs) for idx, (test_dat, names) in enumerate(valid_loader): img_gen = model.test(test_dat, real_input_flag) img_gen = img_gen.transpose(0, 1, 3, 4, 2) # 输出序列 [N S H W C] img_seq = test_dat.detach().cpu().numpy().transpose( 0, 1, 3, 4, 2) # 整个输入序列 [N S H W C] output_length = configs.predict_length img_gen_length = img_gen.shape[1] img_out = img_gen[:, -output_length:] # MSE per frame for i in range(output_length): x = img_seq[:, i + configs.input_length, :, :, :] gx = img_out[:, i, :, :, :] # 对输出数据clip if is_sunspots: gx = np.clip(gx, a_min=0, a_max=1.0) else: gx = np.clip(gx, a_min=-1., a_max=1.0) mse = np.square(x - gx).mean() img_mse[i] += mse avg_mse += mse if is_sunspots: real_frm = np.uint8(x * 255) pred_frm = np.uint8(gx * 255) else: real_frm = np.uint8(x * 127.5 + 127.5) pred_frm = np.uint8(gx * 127.5 + 127.5) for b in range(configs.batch_size): score, _ = ski_ssim(pred_frm[b], real_frm[b], full=True, multichannel=True) ssim[i] += score # save prediction examples if idx <= configs.num_save_samples or (not configs.is_training): path = os.path.join(res_path, names[0][0]) if not os.path.exists(path): os.mkdir(path) name_list = [_[0] for _ in names] np.savez(os.path.join(res_path, '%s.npz' % names[0][0]), inputs=img_seq[0], preds=img_gen[0], names=name_list) for i in range(configs.input_length + configs.predict_length): name = 'gt' + str(i + 1) + '.png' file_name = os.path.join(path, name) if is_sunspots: img_gt = np.uint8(img_seq[0, i, :, :, :] * 255) else: img_gt = np.uint8(img_seq[0, i, :, :, :] * 127.5 + 127.5) cv2.imwrite(file_name, img_gt) for i in range(img_gen_length): name = 'pd' + str(i + 2) + '.png' file_name = os.path.join(path, name) img_pd = img_gen[0, i, :, :, :] if is_sunspots: img_pd = np.clip(img_pd, a_min=0, a_max=1.0) else: img_pd = np.clip(img_pd, a_min=-1., a_max=1.0) if is_sunspots: img_pd = np.uint8(img_pd * 255) else: img_pd = np.uint8(img_pd * 127.5 + 127.5) cv2.imwrite(file_name, img_pd) avg_mse = avg_mse / test_size img_mse = [_ / test_size for _ in img_mse] print('mse of avg: ' + str(avg_mse)) print('mse of seq: ' + str(img_mse)) ssim = [_ / test_size for _ in ssim] avg_ssim = np.mean(ssim) print('ssim of avg: ' + str(avg_ssim)) print('ssim of seq: ' + str(ssim)) f = open(os.path.join(configs.save_dir, 'Metric.txt'), 'a') writer = csv.writer(f, lineterminator='\n') metric = [epoch, avg_mse] + img_mse + ssim writer.writerow(metric) return {'avg_mse': avg_mse, 'avg_ssim': avg_ssim}
(epoch, ti)), padding=0, normalize=True) # Calculate SSIM and PSNR fakeB = fakeB.data.cpu().numpy() realB = realB.data.cpu().numpy() fakeB = fakeB.transpose( (0, 2, 3, 1)) # N C H W ---> N H W C realB = realB.transpose( (0, 2, 3, 1)) # N C H W ---> N H W C fakeB = np.clip(fakeB, a_min=0.0, a_max=1.0) realB = np.clip(realB, a_min=0.0, a_max=1.0) for _bti in range(fakeB.shape[0]): per_fakeB = fakeB[_bti] per_realB = realB[_bti] test_ssim += ski_ssim(per_fakeB, per_realB, data_range=1, multichannel=True) test_psnr += ski_psnr(per_realB, per_fakeB, data_range=1) test_ite += 1 test_psnr /= test_ite test_ssim /= test_ite print(' Valid PSNR: {:.4f}'.format(test_psnr)) print(' Valid SSIM: {:.4f}'.format(test_ssim)) print('------------------------')