def SRCNN2(
args, image_file
): # CHANGE TO INPUT THE after-resize IMAGE FILE, SO IN THE OUTPUT3, NEED TO STORE THE denoise+resize image
# load the SRCNN weights model
#cudnn.benchmark = True
device = torch.device('cuda: 0' if torch.cuda.is_available() else 'cpu')
model = SRCNN().to(device)
state_dict = model.state_dict()
weights_dir = os.getcwd() + '\\SRCNN_outputs\\x{}\\'.format(
args.SR_scale) #
weights_file = os.path.join(weights_dir, 'best.pth') ###
if not weights_file:
print(weights_file + ' not exist')
for n, p in torch.load(weights_file,
map_location=lambda storage, loc: storage).items():
if n in state_dict.keys():
state_dict[n].copy_(p)
else:
raise KeyError(n)
model.eval() # model set in evaluation mode
img_format = image_file[-4:]
image = pil_image.open(image_file).convert('RGB') # 512
image = np.array(image).astype(np.float32)
ycbcr = convert_rgb_to_ycbcr(image)
y = ycbcr[..., 0]
y /= 255.
y = torch.from_numpy(y).to(device)
y = y.unsqueeze(0).unsqueeze(0)
with torch.no_grad():
preds = model(y).clamp(0.0, 1.0) # output2.size 510
# psnr = calc_psnr(y, preds)
# print('PSNR: {:.2f}'.format(psnr))
preds = preds.mul(255.0).cpu().numpy().squeeze(0).squeeze(
0) # tensor -> np
output = np.array([preds, ycbcr[..., 1],
ycbcr[..., 2]]).transpose([1, 2, 0]) # why transpose
output = np.clip(convert_ycbcr_to_rgb(output), 0.0, 255.0).astype(np.uint8)
output = pil_image.fromarray(output)
return output ## type in pil_image
def main():
dataloaders = myDataloader()
test_loader = dataloaders.getTestLoader(batch_size)
model = SRCNN().cuda()
model.load_state_dict(torch.load("./result/train/20srcnnParms.pth"))
model.eval()
with torch.no_grad():
for i, (pic, blurPic, index) in enumerate(test_loader):
blurPic = blurPic.cuda()
out = model(blurPic).cpu()
blurPic = blurPic.cpu()
showImages(pic[0])
showImages(blurPic[0])
showImages(out[0])
print(index)
break
def main():
dataloaders = myDataloader()
test_loader = dataloaders.getTestLoader(batch_size)
model = SRCNN().cuda()
model.load_state_dict(
torch.load("./result/train/" + str(epoch) + "srcnnParms.pth"))
model.eval()
with torch.no_grad():
for i, (pic, blurPic, index) in enumerate(test_loader):
pic = pic.cuda()
blurPic = blurPic.cuda()
out = model(blurPic)
res = torch.cosine_similarity(pic, out, dim=1)
res = res[0]
minValue = torch.min(res)
meanValue = res.mean()
output = 1 - ((res >= meanValue) + 0)
plt.figure()
plt.title(index)
plt.imshow(output.cpu(), cmap="gray")
plt.show()
break
cudnn.benchmark = True
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = SRCNN().to(device)
state_dict = model.state_dict()
for n, p in torch.load(args.weights_file,
map_location=lambda storage, loc: storage).items():
if n in state_dict.keys():
state_dict[n].copy_(p)
else:
raise KeyError(n)
# model, optim = torch.load(model.state_dict(), os.path.join(args.weights_file, 'epoch_150.pth'))
model.eval()
image = Image.open(args.image_file).convert('RGB')
resample = image
image_width = (image.width // args.scale) * args.scale
image_height = (image.height // args.scale) * args.scale
image = image.resize((image_width, image_height), resample=Image.BICUBIC)
image = image.resize(
(image.width // args.scale, image.height // args.scale),
resample=Image.BICUBIC)
image = image.resize((image.width * args.scale, image.height * args.scale),
resample=Image.BICUBIC)
image.save(args.image_file.replace('.',
'_bicubic_x{}.'.format(args.scale)))
return vars(parser.parse_args())
if __name__ == '__main__':
args = get_arguments()
weight = args.get('weight')
p = args.get('input')
upscale = args.get('scale')
dirpath = os.path.dirname(p)
input_filename = os.path.basename(p)
gen = SRCNN()
gen.load_state_dict(torch.load(weight))
gen.eval()
img = Image.open(p)
new_size = [int(x * upscale) for x in img.size]
converter = T.Compose([
T.Resize(size=new_size[::-1], interpolation=Image.BICUBIC),
T.ToTensor()
])
with torch.no_grad():
x = converter(img)
pred = gen(x[None, :, :, :])[0]
save_image(pred, os.path.join(dirpath, f'outputx{upscale}_{input_filename}'))
img = Image.open(path).convert('YCbCr')
y, cb, cr = img.split()
img = img.resize(
(int(img.size[0] * zoom_factor), int(img.size[1] * zoom_factor)),
Image.BICUBIC) # first, we upscale the image via bicubic interpolation
img.save(f'{path.stem}_bicubic.jpg')
img_to_tensor = transforms.ToTensor()
input = img_to_tensor(y).view(
1, -1, y.size[1], y.size[0]) # we only work with the "Y" channel
device = torch.device("cpu")
model = SRCNN()
model.load_state_dict(torch.load('model_21.pth'))
model = model.eval().to(device)
input = input.to(device)
with torch.no_grad():
out = model(input)
out = out.cpu()
out_img_y = out[0].numpy()
out_img_y *= 255.0
out_img_y = out_img_y.clip(0, 255)
out_img_y = Image.fromarray(np.uint8(out_img_y[0]), mode='L')
out_img = Image.merge(
'YCbCr',
[out_img_y,
cb.resize(out_img_y.size),
