Skip to content
GitLab
Commit 333027e0 authored by Marie Tahon's avatar Marie Tahon
Browse files

remove unsed folders

parent 261190ad
Hide whitespace changes
Inline Side-by-side
main_holo.py
View file @ 333027e0
import torch
import time
from utils import *
import nntools as nt
from model import *
from data import *
from argument import *
def save_clean_pred_rad(args, exp, clean_pred_rad, noisy, clean, nom_img = "NoisyPhasePred"):
import os
import torch
from holodncnn.utils import *
from holodncnn.model import DnCNN
from holodncnn.nntools import Experiment, DenoisingStatsManager
from holodncnn.holosets import TrainHoloset, EvalHoloset
import yaml
import argparse
def config():
'''Add arguments
'''
parser = argparse.ArgumentParser(
description='DnCNN')
parser.add_argument('--config', dest='config', type=str, default='./dataset.yaml',
help='configuration file')
parser.add_argument('--test_mode', dest='test_mode', action='store_true',
help='testing phase')
return parser.parse_args()
def save_clean_pred_rad(output_dir, exp, clean_pred_rad, noisy, clean, nom_img="NoisyPhasePred"):
"""This method is used to save the result of a de-noising operation
Arguments:
......@@ -22,18 +32,17 @@ def save_clean_pred_rad(args, exp, clean_pred_rad, noisy, clean, nom_img = "Nois
nom_img (str, optional) : The saving name for the result
"""
save_name = os.path.join(args.save_test_dir, args.input_dir, "Test")
#save_name = os.path.join(args.save_test_dir, args.input_dir, "Test")
save_name = os.path.join(output_dir, "Test")
if not os.path.exists(save_name):
os.makedirs(save_name)
save_images(os.path.join(save_name, '%s-noisy.tiff' % (nom_img)), noisy)
save_images(os.path.join(save_name, '%s-clean.tiff' % (nom_img)), clean)
save_images(os.path.join(save_name , '%s-noisy.tiff' % (nom_img)), noisy)
save_images(os.path.join(save_name , '%s-clean.tiff' % (nom_img)), clean)
save_images(os.path.join(save_name , '%s-%d.tiff' % (nom_img, exp.epoch)), clean_pred_rad)
save_MAT_images(os.path.join(save_name , '%s-%d.mat' % (nom_img, exp.epoch)), clean_pred_rad)
save_images(os.path.join(save_name, '%s-%d.tiff' % (nom_img, exp.epoch)), clean_pred_rad)
save_MAT_images(os.path.join(save_name, '%s-%d.mat' % (nom_img, exp.epoch)), clean_pred_rad)
epoch = exp.epoch
psnr = cal_psnr(rad_to_flat(clean_pred_rad), rad_to_flat(clean))
......@@ -46,7 +55,7 @@ def save_clean_pred_rad(args, exp, clean_pred_rad, noisy, clean, nom_img = "Nois
print("std : ", std)
print("\n")
with open(os.path.join(save_name , '%s-%d.res' % (nom_img, exp.epoch)), 'w') as f:
with open(os.path.join(save_name, '%s-%d.res' % (nom_img, exp.epoch)), 'w') as f:
print("image : ", nom_img, file=f)
print("epoch : ", epoch, file=f)
print("psnr : ", psnr, file=f)
......@@ -55,91 +64,60 @@ def save_clean_pred_rad(args, exp, clean_pred_rad, noisy, clean, nom_img = "Nois
def evaluate_on_HOLODEEP(args, exp):
"""This method is used to run an evaluation on the training database
Arguments:
args (ArgumentParser) : The different info used to do and save the de-noising operations
exp (Experiment) : The de-noising model
def evaluate_with_ref(param, data, exp):
"""Denoise a liste of images and compute metrics when the clean reference is known
- clean is the clean reference of size ()
- noisy is the noisy image to be processed
"""
patterns = args.test_patterns
noises = args.test_noises
clean, noisy = from_DATABASE(args.eval_dir, noises, patterns, True)
clean = np.array(clean)
noisy = np.array(noisy)
running_std = 0
for i in range(noisy.shape[0]):
clean_pred_rad = denoise_img(args, noisy[i], clean[i], "test-{:0>2}".format(i), exp)
std = cal_std_phase(clean_pred_rad, clean[i])
assert (len(data) > 0), "No data to evaluate"
for i in range(len(data)):
noisy, clean = data[i]
clean_pred_rad = noisy
for j in range(param['test']['nb_iteration']):
clean_pred_rad = denoising_single_image(clean_pred_rad, exp)
save_clean_pred_rad(param['test']['save_test_dir'],
exp,
clean_pred_rad.numpy().squeeze(),
noisy.numpy().squeeze(),
clean.numpy().squeeze(),
nom_img="test-{:0>2}".format(i))
std = cal_std_phase(clean_pred_rad, clean)
running_std += std
print("average_std : ", running_std / noisy.shape[0])
print("On the patterns : ", patterns)
print("With noise : ", noises)
print("average_std : ", running_std/noisy.shape[0])
def evaluate_on_DATAEVAL(args, exp):
"""This method is used to run an evaluation on the three test images
Arguments:
args (ArgumentParser) : The different info used to do and save the de-noising operations
exp (Experiment) : The model used to do the de-noising operation
"""
dir_name = args.test_dir
#nameList = ["DATA_1_Phase_Type1_2_0.25_1.5_4_50.mat", "DATA_20_Phase_Type4_2_0.25_2.5_4_100.mat", "VibPhaseDATA.mat"]
nameList = get_files(pathlib.Path(dir_name), '.*.mat')
dataList = []
for name in nameList:
dataList.append(( load_test_data(name, key = "Phaseb", flipupdown = True),
load_test_data(name, key = "Phase", flipupdown = True)))
for idx, (noisy, clean) in enumerate(dataList):
denoise_img(args, noisy, clean, os.path.basename(nameList[idx]), exp)
def denoise_img(args, noisy, clean, name, exp):
"""This method is used to do and save a de-noising operation on a given image
Arguments:
args (ArgumentParser) : The different info used to do and save the de-noising operations
noisy (numpy.array) : The image to de-noise
clean (numpy.array) : The clean reference
name (str) : The name used to save the results
exp (Experiment) : The model used to do the de-noising operation
def evaluate_without_ref(param, data, exp):
"""Denoise a liste of images and compute metrics when the clean reference is known
- clean is the clean reference of size ()
- noisy is the noisy image to be processed
"""
clean_pred_rad = noisy
nb_iteration = args.nb_iteration
for j in range(nb_iteration):
clean_pred_rad = denoising_single_image(args, clean_pred_rad, exp)
save_clean_pred_rad(args, exp, clean_pred_rad, noisy, clean, nom_img = name)
running_std = 0
for i in range(len(data)):
noisy, _ = data[i]
clean_pred_rad = noisy
for j in range(param['test']['nb_iteration']):
clean_pred_rad = denoising_single_image(clean_pred_rad, exp)
save_clean_pred_rad(param['test']['save_test_dir'],
exp,
clean_pred_rad.numpy().squeeze(),
noisy.numpy().squeeze(),
noisy.numpy().squeeze(),
nom_img="test-{:0>2}".format(i))
std = cal_std_phase(clean_pred_rad, noisy)
running_std += std
return clean_pred_rad
print("average_std : ", running_std / noisy.shape[0])
def denoising_single_image(args, noisy, exp):
def denoising_single_image(noisy, exp):
"""This method is used to do a de-noising operation on a given image
Arguments:
......@@ -148,14 +126,16 @@ def denoising_single_image(args, noisy, exp):
exp (Experiment) : The model used to denoise
"""
noisyPy = noisy.reshape(1, args.image_mode, args.test_image_size[0], args.test_image_size[1])
#noisyPy = noisy.reshape(1, args.image_mode, args.test_image_size[0], args.test_image_size[1])
noisyPy_cos = torch.Tensor(normalize_data(noisyPy, 'cos', None))
noisyPy_sin = torch.Tensor(normalize_data(noisyPy, 'sin', None))
# noisyPy_cos = torch.Tensor(normalize_data(noisyPy, 'cos', None))
# noisyPy_sin = torch.Tensor(normalize_data(noisyPy, 'sin', None))
#clean_pred_cos = exp.test(noisyPy_cos).detach().cpu().numpy()
#clean_pred_sin = exp.test(noisyPy_sin).detach().cpu().numpy()
noisyPy_cos = torch.cos(noisy).unsqueeze(0)
noisyPy_sin = torch.sin(noisy).unsqueeze(0)
# clean_pred_cos = exp.test(noisyPy_cos).detach().cpu().numpy()
# clean_pred_sin = exp.test(noisyPy_sin).detach().cpu().numpy()
clean_pred_cos = exp.test(noisyPy_cos)
clean_pred_sin = exp.test(noisyPy_sin)
......@@ -169,70 +149,89 @@ def denoising_single_image(args, noisy, exp):
clean_pred_cos = clean_pred_cos.numpy()
clean_pred_sin = clean_pred_sin.numpy()
clean_pred_rad = np.angle(clean_pred_cos + clean_pred_sin * 1J)
clean_pred_rad = clean_pred_rad.reshape(1, args.test_image_size[0], args.test_image_size[1], args.image_mode)
clean_pred_rad = torch.Tensor(np.angle(clean_pred_cos + clean_pred_sin * 1J)).squeeze(0)
#clean_pred_rad = clean_pred_rad.reshape(1, noisyPy_cos.size(2), noisyPy_cos.size(2), 1)
return clean_pred_rad
def run(args):
def run(param, test_mode):
"""This method is the main method
"""
device = 'cuda' if torch.cuda.is_available() else 'cpu'
net = DnCNN(D=args.D, C=args.C, image_mode=args.image_mode).to(device)
adam = torch.optim.Adam(net.parameters(), lr=args.lr)
statsManager = DenoisingStatsManager()
exp = nt.Experiment(net, adam, statsManager, perform_validation_during_training=args.perform_validation, input_dir=args.input_dir, output_dir=args.output_dir, startEpoch=args.epoch, freq_save=args.freq_save)
if not args.test_mode :
print("\n=>Training until epoch :<===\n", args.num_epochs)
net = DnCNN(D=param['model']['D'],
C=param['model']['C'],
image_mode=1
).to(device)
adam = torch.optim.Adam(net.parameters(), lr=param['model']['lr'])
exp = Experiment(net,
adam,
DenoisingStatsManager(),
perform_validation_during_training=param['model']['perform_validation'],
input_dir=param['model']['input_dir'],
output_dir=param['model']['output_dir'],
startEpoch=param['model']['start_epoch'],
freq_save=param['model']['freq_save'])
if not test_mode:
print("\n=>Training until epoch :<===", param['model']['num_epochs'])
print("\n\Model training")
trainData = TrainDataset(args.clean_train, args.noisy_train, args.image_mode, args.train_image_size, nb_rotation=args.nb_rotation)
evalData = EvalDataset(args.eval_dir, args.eval_noises, args.eval_patterns, args.image_mode, args.eval_image_size)
exp.initData(trainData, evalData, batch_size=args.batch_size)
exp.run(num_epochs=args.num_epochs)
if(args.graph):
trainData = TrainHoloset(param['train']['path'],
param['train']['csv'],
param['train']['extension'],
param['train']['matlab_key_clean'],
param['train']['matlab_key_noisy'],
param['train']['augmentation'],
param['train']['patch']['nb_patch_per_image'],
param['train']['patch']['size'],
param['train']['patch']['stride'],
param['train']['patch']['step'],
)
evalData = EvalHoloset(param['eval']['path'],
param['eval']['csv'],
param['eval']['extension'],
param['eval']['matlab_key_clean'],
param['eval']['matlab_key_noisy'],
)
exp.initData(trainData, evalData, batch_size=param['model']['batch_size'])
exp.run(num_epochs=param['model']['num_epochs'])
if (param['model']['graph']):
exp.trace()
else :
print("args.noisy_img : ", args.noisy_img)
if args.noisy_img is None:
evaluate_on_HOLODEEP(args, exp)
evaluate_on_DATAEVAL(args, exp)
else :
noisy = load_test_data(args.noisy_img, key = args.noisy_key, flipupdown=args.flip)
clean = load_test_data(args.clean_img, key = args.clean_key, flipupdown=args.flip)
else:
denoise_img(args, noisy, clean, os.path.basename(args.noisy_img), exp)
testData = EvalHoloset(param['test']['path'],
param['test']['csv'],
param['test']['extension'],
param['test']['matlab_key_clean'],
param['test']['matlab_key_noisy'],
)
print("Nb of test noisy images : ", len(testData))
print("This test is with reference : ", testData.ref)
if testData.ref:
evaluate_with_ref(param, testData, exp)
else:
evaluate_without_ref(param, testData, exp)
if __name__ == '__main__':
args = config()
with open(args.config, 'r') as f:
param = yaml.load(f, Loader=yaml.FullLoader)
args = parse()
print("\n\n")
list(map(lambda p : print(p + " : ", vars(args)[p]), vars(args)))
print("\n\n")
print("\n")
list(map(lambda p: print(p + " : ", param[p]), param))
print("\n")
torch.manual_seed(123456789)
timeElapsed = time.time()
run(args)
run(param, args.test_mode)
print("Time elapsed : ", time.time() - timeElapsed)
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment