diff --git a/codes/data/__init__.py b/codes/data/__init__.py index 5996436c..d8c2447f 100644 --- a/codes/data/__init__.py +++ b/codes/data/__init__.py @@ -39,6 +39,8 @@ def create_dataset(dataset_opt): from data.combined_dataset import CombinedDataset as D elif mode == 'multiscale': from data.multiscale_dataset import MultiScaleDataset as D + elif mode == 'paired_frame': + from data.paired_frame_dataset import PairedFrameDataset as D else: raise NotImplementedError('Dataset [{:s}] is not recognized.'.format(mode)) dataset = D(dataset_opt) diff --git a/codes/data/paired_frame_dataset.py b/codes/data/paired_frame_dataset.py new file mode 100644 index 00000000..2c1be9ac --- /dev/null +++ b/codes/data/paired_frame_dataset.py @@ -0,0 +1,91 @@ +from data.base_unsupervised_image_dataset import BaseUnsupervisedImageDataset +import numpy as np +import torch +from bisect import bisect_left +import os.path as osp + +class PairedFrameDataset(BaseUnsupervisedImageDataset): + def __init__(self, opt): + super(PairedFrameDataset, self).__init__(opt) + + def get_pair(self, chunk_index, chunk_offset): + imname = osp.basename(self.chunks[chunk_index].path) + if '_left' in imname: + chunks = [chunk_index, chunk_index+1] + else: + chunks = [chunk_index-1, chunk_index] + hqs, refs, masks, centers = [], [], [], [] + for i in chunks: + h, r, c, m, p = self.chunks[i][chunk_offset] + hqs.append(h) + refs.append(r) + masks.append(m) + centers.append(c) + path = p + return hqs, refs, masks, centers, path + + def __getitem__(self, item): + chunk_ind = bisect_left(self.starting_indices, item) + chunk_ind = chunk_ind if chunk_ind < len(self.starting_indices) and self.starting_indices[chunk_ind] == item else chunk_ind-1 + hqs, refs, masks, centers, path = self.get_pair(chunk_ind, item-self.starting_indices[chunk_ind]) + + hs, hrs, hms, hcs = self.resize_hq(hqs, refs, masks, centers) + ls, lrs, lms, lcs = self.synthesize_lq(hs, hrs, hms, hcs) + + # Convert to torch tensor + hq = torch.from_numpy(np.ascontiguousarray(np.transpose(np.stack(hs), (0, 3, 1, 2)))).float() + hq_ref = torch.from_numpy(np.ascontiguousarray(np.transpose(np.stack(hrs), (0, 3, 1, 2)))).float() + hq_mask = torch.from_numpy(np.ascontiguousarray(np.stack(hms))).unsqueeze(dim=1) + hq_ref = torch.cat([hq_ref, hq_mask], dim=1) + lq = torch.from_numpy(np.ascontiguousarray(np.transpose(np.stack(ls), (0, 3, 1, 2)))).float() + lq_ref = torch.from_numpy(np.ascontiguousarray(np.transpose(np.stack(lrs), (0, 3, 1, 2)))).float() + lq_mask = torch.from_numpy(np.ascontiguousarray(np.stack(lms))).unsqueeze(dim=1) + lq_ref = torch.cat([lq_ref, lq_mask], dim=1) + + return {'GT_path': path, 'LQ': lq, 'GT': hq, 'gt_fullsize_ref': hq_ref, 'lq_fullsize_ref': lq_ref, + 'lq_center': torch.tensor(lcs, dtype=torch.long), 'gt_center': torch.tensor(hcs, dtype=torch.long)} + + +if __name__ == '__main__': + opt = { + 'name': 'amalgam', + 'paths': ['F:\\4k6k\\datasets\\ns_images\\vr\\paired_images'], + 'weights': [1], + 'target_size': 128, + 'force_multiple': 32, + 'scale': 2, + 'eval': False, + 'fixed_corruptions': ['jpeg-medium'], + 'random_corruptions': [], + 'num_corrupts_per_image': 0, + 'num_frames': 10 + } + + ds = PairedFrameDataset(opt) + import os + os.makedirs("debug", exist_ok=True) + bs = 0 + batch = None + for i in range(len(ds)): + import random + k = 'LQ' + element = ds[random.randint(0,len(ds))] + base_file = osp.basename(element["GT_path"]) + o = element[k].unsqueeze(0) + if bs < 32: + if batch is None: + batch = o + else: + batch = torch.cat([batch, o], dim=0) + bs += 1 + continue + + if 'path' not in k and 'center' not in k: + b, fr, f, h, w = batch.shape + for j in range(fr): + import torchvision + base=osp.basename(base_file) + torchvision.utils.save_image(batch[:, j], "debug/%i_%s_%i__%s.png" % (i, k, j, base)) + + bs = 0 + batch = None