def main():
global opt, model
opt = parser.parse_args()
opt.gpus = '0'
print(opt)
os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus
cuda = opt.cuda
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
opt.seed = random.randint(1, 10000)
torch.manual_seed(opt.seed)
if cuda:
torch.cuda.manual_seed(opt.seed)
cudnn.benchmark = True
print("===> Loading datasets")
train_set = DatasetFromHdf5("./data/training_Gray_5to50_uint8_samples.h5")
training_data_loader = DataLoader(dataset=train_set,
num_workers=opt.threads,
batch_size=opt.batchSize,
shuffle=True,
pin_memory=True)
print("===> Building model")
model = _NetG()
criterion = nn.L1Loss()
print("===> Setting GPU")
if cuda:
model = model.cuda()
criterion = criterion.cuda()
summary(model, (1, 64, 64))
# optionally resume from a checkpoint
if opt.resume:
if os.path.isfile(opt.resume):
print("=> loading checkpoint '{}'".format(opt.resume))
checkpoint = torch.load(opt.resume,
map_location=lambda storage, loc: storage)
opt.start_epoch = checkpoint["epoch"]
model.load_state_dict(checkpoint['model'].state_dict())
del checkpoint
torch.cuda.empty_cache()
else:
print("=> no checkpoint found at '{}'".format(opt.resume))
print("===> Setting Optimizer")
optimizer = optim.Adam(model.parameters(), lr=opt.lr)
print("===> Training")
max_psnr = 0
for epoch in range(opt.start_epoch, opt.nEpochs + 1):
max_psnr = train(training_data_loader, optimizer, model, criterion,
epoch, max_psnr)
save_checkpoint(model, epoch, 'end', 'end_ep')
def output_psnr_mse(img_orig, img_out):
squared_error = np.square(img_orig - img_out)
mse = np.mean(squared_error)
psnr = 10 * np.log10(1.0 / mse)
return psnr
if cuda:
print("=> use gpu id: '{}'".format(opt.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus
if not torch.cuda.is_available():
raise Exception(
"No GPU found or Wrong gpu id, please run without --cuda")
model = _NetG()
checkpoint = torch.load(opt.model, map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint['model'].state_dict())
with torch.no_grad():
model.eval()
if cuda:
model = model.cuda()
else:
model = model.cpu()
origin_list = sorted(
glob.glob("./data/Test/BSD68/original_mat_int/" + "*.*"))
noisy_list = sorted(
glob.glob("./data/Test/BSD68/noisy_mat_s10_int/" + "*.*"))
origin_names = os.listdir("./data/Test/BSD68/original_mat_int/")