TaoHuang2018
commented
2 years ago I am sorry for my late reply. For the training of Neighbor2Neighbor in the SIDD dataset:
- The code for data preparation is in the newly uploaded file
dataset_tool_raw.py. It would be better to consider the camera device difference in the SIDD Medium Set. However, as the device info is not provided in the SIDD validation set, we simply ignore the device difference for training set preparation. - The training code for the SIDD dataset is similar to that for the ImageNet validation set, except some operations for raw images. To mention, the training scheme for the SIDD dataset is also similar, i.e., the number of epoch is 100 and the learning rate decay schedule is the same, except we use a smaller $gamma=1$. Here is some code for processing raw images.
def space_to_depth(x, block_size): n, c, h, w = x.size() unfolded_x = torch.nn.functional.unfold(x, block_size, stride=block_size) return unfolded_x.view(n, c * block_size**2, h // block_size, w // block_size)
def depth_to_space(x, block_size): return torch.nn.functional.pixel_shuffle(x, block_size)
def generate_mask_pair(img):
prepare masks (N x C x H/2 x W/2)
n, c, h, w = img.shape
mask1 = torch.zeros(size=(n * h // 2 * w // 2 * 4, ),
dtype=torch.bool,
device=img.device)
mask2 = torch.zeros(size=(n * h // 2 * w // 2 * 4, ),
dtype=torch.bool,
device=img.device)
# prepare random mask pairs
idx_pair = torch.tensor(
[[0, 1], [0, 2], [1, 3], [2, 3], [1, 0], [2, 0], [3, 1], [3, 2]],
dtype=torch.int64,
device=img.device)
rd_idx = torch.zeros(size=(n * h // 2 * w // 2, ),
dtype=torch.int64,
device=img.device)
torch.randint(low=0,
high=8,
size=(n * h // 2 * w // 2, ),
generator=get_generator(),
out=rd_idx)
rd_pair_idx = idx_pair[rd_idx]
rd_pair_idx += torch.arange(start=0,
end=n * h // 2 * w // 2 * 4,
step=4,
dtype=torch.int64,
device=img.device).reshape(-1, 1)
# get masks
mask1[rd_pair_idx[:, 0]] = 1
mask2[rd_pair_idx[:, 1]] = 1
return mask1, mask2def generate_subimages(img, mask): n, c, h, w = img.shape subimage = torch.zeros(n, c, h // 2, w // 2, dtype=img.dtype, layout=img.layout, device=img.device)
per channel
for i in range(c):
img_per_channel = space_to_depth(img[:, i:i + 1, :, :], block_size=2)
img_per_channel = img_per_channel.permute(0, 2, 3, 1).reshape(-1)
subimage[:, i:i + 1, :, :] = img_per_channel[mask].reshape(
n, h // 2, w // 2, 1).permute(0, 3, 1, 2)
return subimageclass DataLoader_SIDD_Medium_Raw(data.Dataset): def init(self, data_dir): super(DataLoader_SIDD_Medium_Raw, self).init() self.data_dir = data_dir
get images path
self.train_fns = glob.glob(os.path.join(self.data_dir, "*"))
self.train_fns.sort()
print('fetch {} samples for training'.format(len(self.train_fns)))
def __getitem__(self, index):
# fetch image
fn = self.train_fns[index]
im = loadmat(fn)["x"]
im = im[np.newaxis, :, :]
im = torch.from_numpy(im)
return im
def __len__(self):
return len(self.train_fns)def get_SIDD_validation(dataset_dir): val_data_dict = loadmat( os.path.join(dataset_dir, "ValidationNoisyBlocksRaw.mat")) val_data_noisy = val_data_dict['ValidationNoisyBlocksRaw'] val_data_dict = loadmat( os.path.join(dataset_dir, 'ValidationGtBlocksRaw.mat')) val_data_gt = val_data_dict['ValidationGtBlocksRaw'] num_img, numblock, , _ = val_data_gt.shape return num_img, num_block, val_data_noisy, val_data_gt
madfff
zejinwang
Thanks for your great work. I apply neighbor sub-sampler on the packed 4-channel raw images and cannot reproduce the results. Did I do anything wrong? Here is my data-processing code.