Open ShiinaMitsuki opened 4 years ago
import torch from torch.nn.functional import conv2d from torch.nn import Conv2d import numpy as np
def srm_filter(x): q = [4.0, 12.0, 2.0] filter1 = [[0, 0, 0, 0, 0], [0, -1, 2, -1, 0], [0, 2, -4, 2, 0], [0, -1, 2, -1, 0], [0, 0, 0, 0, 0]] filter2 = [[-1, 2, -2, 2, -1], [2, -6, 8, -6, 2], [-2, 8, -12, 8, -2], [2, -6, 8, -6, 2], [-1, 2, -2, 2, -1]] filter3 = [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 1, -2, 1, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]] filter1 = np.asarray(filter1, dtype=np.float32) / q[0] filter2 = np.asarray(filter2, dtype=np.float32) / q[1] filter3 = np.asarray(filter3, dtype=np.float32) / q[2] filters = np.asarray( [ [filter3, filter3, filter3], # R [filter2, filter2, filter2], # G [filter1, filter1, filter1], # B ], dtype=np.float32)
output = conv2d(torch.from_numpy(x.astype(np.float32)), torch.from_numpy(filters), stride=1, padding=2)
return output
if name == 'main': from PIL import Image
raw_img = Image.open(
r'./2.png').convert('RGB')
noisemap = srm_filter((np.asarray(raw_img) / 255.0).transpose(2, 0, 1)[np.newaxis, ...])
noisemap = Image.fromarray((noisemap.squeeze(dim=0).numpy() * 255.0).transpose(1, 2, 0).astype(np.uint8))
noisemap.save('noise_map.png')
raw_img.save('raw_map.png')
![Uploading 2.png…]()
Sorry but I can't find the implementation code of the noise input using SRM filter?
I write one using pytorch but the output seem different from paper, here's the code.
and here's the output:
the output from paper: