25 lines
878 B
Python
25 lines
878 B
Python
import os
|
|
import numpy as np
|
|
from models import Noise2Same
|
|
import tensorflow as tf
|
|
import random
|
|
|
|
os.environ['PYTHONHASHSEED'] = '1'
|
|
random.seed(666)
|
|
np.random.seed(666)
|
|
tf.set_random_seed(666)
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = '0' # Adjust to choose GPU you want to use
|
|
|
|
data_dir = 'Denoising_data/'
|
|
X = np.load(data_dir+'train/DCNN400_train_gaussian25.npy')
|
|
X_val = np.load(data_dir+'val/DCNN400_validation_gaussian25.npy')
|
|
X = np.array([(x - x.mean())/x.std() for x in X])
|
|
X_val = np.array([(x - x.mean())/x.std() for x in X_val]).astype('float32')
|
|
|
|
model_dir = 'N2S-3000' # Set model checkpoints save path
|
|
steps = 3000 # Set training steps
|
|
|
|
sgm_loss = 1 # the default sigma is 1
|
|
model = Noise2Same('trained_models/', model_dir, dim=2, in_channels=1, lmbd=2*sgm_loss)
|
|
model.train(X[..., None], patch_size=[64, 64], validation=X_val[..., None], batch_size=64, steps=steps)
|