From 8e5b6682bf055b2c971d7161ac7b81ae75d9ef34 Mon Sep 17 00:00:00 2001
From: James Betker <jbetker@gmail.com>
Date: Fri, 23 Oct 2020 20:58:07 -0600
Subject: [PATCH] Add PairedFrameDataset

---
 codes/data/__init__.py             |  2 +
 codes/data/paired_frame_dataset.py | 91 ++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+)
 create mode 100644 codes/data/paired_frame_dataset.py

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