如何使ESRGAN与.raw图像配合使用?

时间:2019-05-05 18:37:50

标签: python image numpy pytorch

我正在对 .raw 图像进行缩放,并希望使用ESRGAN,因为众所周知它可以与 .png 一起使用。如何添加 .raw 支持?

我尝试更改图像处理代码,因为 .raw 的规范与 .png 不同。 预训练的模型适用于RGB图像,因此需要进行调整。我尝试遵循此guide,但是它显示 AssertionError C:/ Users / HP / Desktop / BasicSR-master / datasets / woodhr /没有有效的图像文件

import sys
import numpy as np
import torch
import architecture as arch

model_path = sys.argv[1]  # models/RRDB_ESRGAN_x4.pth OR models/RRDB_PSNR_x4.pth
device = torch.device('cpu')  # if you want to run on CPU, change 'cuda' -> cpu

test_img_folder = 'LR/*'

model = arch.RRDB_Net(3, 3, 64, 23, gc=32, upscale=4, norm_type=None, act_type='leakyrelu', \
                      mode='CNA', res_scale=1, upsample_mode='upconv')
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
for k, v in model.named_parameters():
    v.requires_grad = False
model = model.to(device)

print('Model path {:s}. \nTesting...'.format(model_path))

fd = open('LR/N47E006.raw', 'rb')

rows, cols = (257, 257)
fd = open("C:\\Users\\HP\\Desktop\\ESRGAN\LR\\N47E006_257.raw", "rb")
f = np.fromfile(fd, dtype=np.uint16, count=rows * cols)
img = f.reshape((rows, cols))
img = np.broadcast_to(img[..., np.newaxis], (257, 257, 3))
fd.close()

darr = np.array(img)
minimum = darr.min()
maximum = darr.max()

img = (img - minimum) * (1 - 0) / (maximum - minimum) + 0

img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
img_LR = img.unsqueeze(0)
img_LR = img_LR.to(device)

output = model(img_LR).data.squeeze().float().cpu().clamp_(0, 1).numpy()
output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))

output = (output * (maximum - minimum) + minimum).round()

l = output.astype('int')

l.astype('uint16').tofile('results/N47E006_257_rlt.raw')

我希望输出与输入类似,但更大。输入的是带有梯级的高度图。但是,实际输出是具有垂直线的平面。 总之,我看到了两个任务。

  1. 编写适合 .raw 图片的代码。
  2. 调整 .raw 的模型。

0 个答案:

没有答案