19 lines
733 B
Python
19 lines
733 B
Python
import os
|
|
import numpy as np
|
|
from models import Noise2Same
|
|
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = '0' # Adjust to choose GPU you want to use
|
|
|
|
def PSNR(gt, img):
|
|
mse = np.mean(np.square(gt - img))
|
|
return 20 * np.log10(255) - 10 * np.log10(mse)
|
|
|
|
model_dir = 'N2S-3000' # Adjust your model path
|
|
data_dir = 'Denoising_data/test/'
|
|
model = Noise2Same('trained_models/', model_dir, dim=2, in_channels=1)
|
|
groundtruth_data = np.load(data_dir+'bsd68_groundtruth.npy', allow_pickle=True)
|
|
test_data = np.load(data_dir+'bsd68_gaussian25.npy', allow_pickle=True)
|
|
preds = [model.predict(d.astype('float32')) for d in test_data]
|
|
psnrs = [PSNR(preds[idx], groundtruth_data[idx]) for idx in range(len(test_data))]
|
|
print(np.array(psnrs).mean())
|