def predict(self, img_path):
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
SR_dir = join(os.path.dirname(img_path), 'result')
if not os.path.exists(SR_dir):
os.makedirs(SR_dir)
testloader = DataLoader(DataValSet(img_path),
batch_size=1,
shuffle=False,
pin_memory=False)
model = torch.load(self.model)
model = model.to(device)
with torch.no_grad():
for iteration, batch in enumerate(testloader, 1):
LR_Blur = batch[0]
HR = batch[1]
LR_Blur = LR_Blur.to(device)
HR = HR.to(device)
test_Tensor = torch.cuda.FloatTensor().resize_(1).zero_() + 1
gated_Tensor = torch.cuda.FloatTensor().resize_(1).zero_() + 1
[lr_deblur, sr] = model(LR_Blur, gated_Tensor, test_Tensor)
sr = torch.clamp(sr, min=0, max=1)
torch.cuda.synchronize() #wait for CPU & GPU time syn
resultSRDeblur = transforms.ToPILImage()(sr.cpu()[0])
resultSRDeblur.save(join(SR_dir, 'result.png'))
criterion = torch.nn.MSELoss(size_average=True)
criterion = criterion.to(device)
print(opt)
test(testloader, model, criterion, SR_dir)
opt = parser.parse_args()
root_val_dir = opt.dataset # #----------Validation path
SR_dir = join(root_val_dir,
'Results') #--------------------------SR results save path
isexists = os.path.exists(SR_dir)
if not isexists:
os.makedirs(SR_dir)
print("The results of testing images sotre in {}.".format(SR_dir))
testloader = DataLoader(DataValSet(root_val_dir),
batch_size=1,
shuffle=False,
pin_memory=False)
print("===> Loading model and criterion")
if opt.intermediate_process:
test_pkl = opt.intermediate_process
if is_pkl(test_pkl):
print("Testing model {}----------------------------------".format(
opt.intermediate_process))
train_step, epoch = which_trainingstep_epoch(opt.intermediate_process)
opt.gated = test_set[train_step - 1]['gated']
model = torch.load(test_pkl)
model_test(model)
else:
def model_test(model):
model = model.to(device)
criterion = torch.nn.MSELoss(size_average=True)
criterion = criterion.to(device)
print(opt)
test(testloader, model, criterion, SR_dir)
opt = parser.parse_args()
root_val_dir = opt.dataset# #----------Validation path
SR_dir = join(root_val_dir, 'Results') #--------------------------SR results save path
isexists = os.path.exists(SR_dir)
if not isexists:
os.makedirs(SR_dir)
print("The results of testing images sotre in {}.".format(SR_dir))
testloader = DataLoader(DataValSet(root_val_dir), batch_size=1, shuffle=False, pin_memory=False)
print("===> Loading model and criterion")
if opt.intermediate_process:
test_pkl = opt.intermediate_process
if is_pkl(test_pkl):
print("Testing model {}----------------------------------".format(opt.intermediate_process))
train_step, epoch = which_trainingstep_epoch(opt.intermediate_process)
opt.gated = test_set[train_step-1]['gated']
model = torch.load(test_pkl)
model_test(model)
else:
print("It's not a pkl file. Please give a correct pkl folder on command line for example --opt.intermediate_process /models/1/GFN_epoch_25.pkl)")
else:
test_dir = 'models/'
test_list = [x for x in sorted(os.listdir(test_dir)) if is_pkl(x)]