Add files via upload

1 files changed, 360 insertions, 0 deletions

diff --git a/r_basicsr/data/reds_dataset.py b/r_basicsr/data/reds_dataset.py
new file mode 100644
index 0000000..fa7df26
--- /dev/null
+++ b/r_basicsr/data/reds_dataset.py
@@ -0,0 +1,360 @@
+import numpy as np

+import random

+import torch

+from pathlib import Path

+from torch.utils import data as data

+

+from r_basicsr.data.transforms import augment, paired_random_crop

+from r_basicsr.utils import FileClient, get_root_logger, imfrombytes, img2tensor

+from r_basicsr.utils.flow_util import dequantize_flow

+from r_basicsr.utils.registry import DATASET_REGISTRY

+

+

+@DATASET_REGISTRY.register()

+class REDSDataset(data.Dataset):

+    """REDS dataset for training.

+

+    The keys are generated from a meta info txt file.

+    basicsr/data/meta_info/meta_info_REDS_GT.txt

+

+    Each line contains:

+    1. subfolder (clip) name; 2. frame number; 3. image shape, separated by

+    a white space.

+    Examples:

+    000 100 (720,1280,3)

+    001 100 (720,1280,3)

+    ...

+

+    Key examples: "000/00000000"

+    GT (gt): Ground-Truth;

+    LQ (lq): Low-Quality, e.g., low-resolution/blurry/noisy/compressed frames.

+

+    Args:

+        opt (dict): Config for train dataset. It contains the following keys:

+            dataroot_gt (str): Data root path for gt.

+            dataroot_lq (str): Data root path for lq.

+            dataroot_flow (str, optional): Data root path for flow.

+            meta_info_file (str): Path for meta information file.

+            val_partition (str): Validation partition types. 'REDS4' or

+                'official'.

+            io_backend (dict): IO backend type and other kwarg.

+

+            num_frame (int): Window size for input frames.

+            gt_size (int): Cropped patched size for gt patches.

+            interval_list (list): Interval list for temporal augmentation.

+            random_reverse (bool): Random reverse input frames.

+            use_hflip (bool): Use horizontal flips.

+            use_rot (bool): Use rotation (use vertical flip and transposing h

+                and w for implementation).

+

+            scale (bool): Scale, which will be added automatically.

+    """

+

+    def __init__(self, opt):

+        super(REDSDataset, self).__init__()

+        self.opt = opt

+        self.gt_root, self.lq_root = Path(opt['dataroot_gt']), Path(opt['dataroot_lq'])

+        self.flow_root = Path(opt['dataroot_flow']) if opt['dataroot_flow'] is not None else None

+        assert opt['num_frame'] % 2 == 1, (f'num_frame should be odd number, but got {opt["num_frame"]}')

+        self.num_frame = opt['num_frame']

+        self.num_half_frames = opt['num_frame'] // 2

+

+        self.keys = []

+        with open(opt['meta_info_file'], 'r') as fin:

+            for line in fin:

+                folder, frame_num, _ = line.split(' ')

+                self.keys.extend([f'{folder}/{i:08d}' for i in range(int(frame_num))])

+

+        # remove the video clips used in validation

+        if opt['val_partition'] == 'REDS4':

+            val_partition = ['000', '011', '015', '020']

+        elif opt['val_partition'] == 'official':

+            val_partition = [f'{v:03d}' for v in range(240, 270)]

+        else:

+            raise ValueError(f'Wrong validation partition {opt["val_partition"]}.'

+                             f"Supported ones are ['official', 'REDS4'].")

+        self.keys = [v for v in self.keys if v.split('/')[0] not in val_partition]

+

+        # file client (io backend)

+        self.file_client = None

+        self.io_backend_opt = opt['io_backend']

+        self.is_lmdb = False

+        if self.io_backend_opt['type'] == 'lmdb':

+            self.is_lmdb = True

+            if self.flow_root is not None:

+                self.io_backend_opt['db_paths'] = [self.lq_root, self.gt_root, self.flow_root]

+                self.io_backend_opt['client_keys'] = ['lq', 'gt', 'flow']

+            else:

+                self.io_backend_opt['db_paths'] = [self.lq_root, self.gt_root]

+                self.io_backend_opt['client_keys'] = ['lq', 'gt']

+

+        # temporal augmentation configs

+        self.interval_list = opt['interval_list']

+        self.random_reverse = opt['random_reverse']

+        interval_str = ','.join(str(x) for x in opt['interval_list'])

+        logger = get_root_logger()

+        logger.info(f'Temporal augmentation interval list: [{interval_str}]; '

+                    f'random reverse is {self.random_reverse}.')

+

+    def __getitem__(self, index):

+        if self.file_client is None:

+            self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)

+

+        scale = self.opt['scale']

+        gt_size = self.opt['gt_size']

+        key = self.keys[index]

+        clip_name, frame_name = key.split('/')  # key example: 000/00000000

+        center_frame_idx = int(frame_name)

+

+        # determine the neighboring frames

+        interval = random.choice(self.interval_list)

+

+        # ensure not exceeding the borders

+        start_frame_idx = center_frame_idx - self.num_half_frames * interval

+        end_frame_idx = center_frame_idx + self.num_half_frames * interval

+        # each clip has 100 frames starting from 0 to 99

+        while (start_frame_idx < 0) or (end_frame_idx > 99):

+            center_frame_idx = random.randint(0, 99)

+            start_frame_idx = (center_frame_idx - self.num_half_frames * interval)

+            end_frame_idx = center_frame_idx + self.num_half_frames * interval

+        frame_name = f'{center_frame_idx:08d}'

+        neighbor_list = list(range(start_frame_idx, end_frame_idx + 1, interval))

+        # random reverse

+        if self.random_reverse and random.random() < 0.5:

+            neighbor_list.reverse()

+

+        assert len(neighbor_list) == self.num_frame, (f'Wrong length of neighbor list: {len(neighbor_list)}')

+

+        # get the GT frame (as the center frame)

+        if self.is_lmdb:

+            img_gt_path = f'{clip_name}/{frame_name}'

+        else:

+            img_gt_path = self.gt_root / clip_name / f'{frame_name}.png'

+        img_bytes = self.file_client.get(img_gt_path, 'gt')

+        img_gt = imfrombytes(img_bytes, float32=True)

+

+        # get the neighboring LQ frames

+        img_lqs = []

+        for neighbor in neighbor_list:

+            if self.is_lmdb:

+                img_lq_path = f'{clip_name}/{neighbor:08d}'

+            else:

+                img_lq_path = self.lq_root / clip_name / f'{neighbor:08d}.png'

+            img_bytes = self.file_client.get(img_lq_path, 'lq')

+            img_lq = imfrombytes(img_bytes, float32=True)

+            img_lqs.append(img_lq)

+

+        # get flows

+        if self.flow_root is not None:

+            img_flows = []

+            # read previous flows

+            for i in range(self.num_half_frames, 0, -1):

+                if self.is_lmdb:

+                    flow_path = f'{clip_name}/{frame_name}_p{i}'

+                else:

+                    flow_path = (self.flow_root / clip_name / f'{frame_name}_p{i}.png')

+                img_bytes = self.file_client.get(flow_path, 'flow')

+                cat_flow = imfrombytes(img_bytes, flag='grayscale', float32=False)  # uint8, [0, 255]

+                dx, dy = np.split(cat_flow, 2, axis=0)

+                flow = dequantize_flow(dx, dy, max_val=20, denorm=False)  # we use max_val 20 here.

+                img_flows.append(flow)

+            # read next flows

+            for i in range(1, self.num_half_frames + 1):

+                if self.is_lmdb:

+                    flow_path = f'{clip_name}/{frame_name}_n{i}'

+                else:

+                    flow_path = (self.flow_root / clip_name / f'{frame_name}_n{i}.png')

+                img_bytes = self.file_client.get(flow_path, 'flow')

+                cat_flow = imfrombytes(img_bytes, flag='grayscale', float32=False)  # uint8, [0, 255]

+                dx, dy = np.split(cat_flow, 2, axis=0)

+                flow = dequantize_flow(dx, dy, max_val=20, denorm=False)  # we use max_val 20 here.

+                img_flows.append(flow)

+

+            # for random crop, here, img_flows and img_lqs have the same

+            # spatial size

+            img_lqs.extend(img_flows)

+

+        # randomly crop

+        img_gt, img_lqs = paired_random_crop(img_gt, img_lqs, gt_size, scale, img_gt_path)

+        if self.flow_root is not None:

+            img_lqs, img_flows = img_lqs[:self.num_frame], img_lqs[self.num_frame:]

+

+        # augmentation - flip, rotate

+        img_lqs.append(img_gt)

+        if self.flow_root is not None:

+            img_results, img_flows = augment(img_lqs, self.opt['use_hflip'], self.opt['use_rot'], img_flows)

+        else:

+            img_results = augment(img_lqs, self.opt['use_hflip'], self.opt['use_rot'])

+

+        img_results = img2tensor(img_results)

+        img_lqs = torch.stack(img_results[0:-1], dim=0)

+        img_gt = img_results[-1]

+

+        if self.flow_root is not None:

+            img_flows = img2tensor(img_flows)

+            # add the zero center flow

+            img_flows.insert(self.num_half_frames, torch.zeros_like(img_flows[0]))

+            img_flows = torch.stack(img_flows, dim=0)

+

+        # img_lqs: (t, c, h, w)

+        # img_flows: (t, 2, h, w)

+        # img_gt: (c, h, w)

+        # key: str

+        if self.flow_root is not None:

+            return {'lq': img_lqs, 'flow': img_flows, 'gt': img_gt, 'key': key}

+        else:

+            return {'lq': img_lqs, 'gt': img_gt, 'key': key}

+

+    def __len__(self):

+        return len(self.keys)

+

+

+@DATASET_REGISTRY.register()

+class REDSRecurrentDataset(data.Dataset):

+    """REDS dataset for training recurrent networks.

+

+    The keys are generated from a meta info txt file.

+    basicsr/data/meta_info/meta_info_REDS_GT.txt

+

+    Each line contains:

+    1. subfolder (clip) name; 2. frame number; 3. image shape, separated by

+    a white space.

+    Examples:

+    000 100 (720,1280,3)

+    001 100 (720,1280,3)

+    ...

+

+    Key examples: "000/00000000"

+    GT (gt): Ground-Truth;

+    LQ (lq): Low-Quality, e.g., low-resolution/blurry/noisy/compressed frames.

+

+    Args:

+        opt (dict): Config for train dataset. It contains the following keys:

+            dataroot_gt (str): Data root path for gt.

+            dataroot_lq (str): Data root path for lq.

+            dataroot_flow (str, optional): Data root path for flow.

+            meta_info_file (str): Path for meta information file.

+            val_partition (str): Validation partition types. 'REDS4' or

+                'official'.

+            io_backend (dict): IO backend type and other kwarg.

+

+            num_frame (int): Window size for input frames.

+            gt_size (int): Cropped patched size for gt patches.

+            interval_list (list): Interval list for temporal augmentation.

+            random_reverse (bool): Random reverse input frames.

+            use_hflip (bool): Use horizontal flips.

+            use_rot (bool): Use rotation (use vertical flip and transposing h

+                and w for implementation).

+

+            scale (bool): Scale, which will be added automatically.

+    """

+

+    def __init__(self, opt):

+        super(REDSRecurrentDataset, self).__init__()

+        self.opt = opt

+        self.gt_root, self.lq_root = Path(opt['dataroot_gt']), Path(opt['dataroot_lq'])

+        self.num_frame = opt['num_frame']

+

+        self.keys = []

+        with open(opt['meta_info_file'], 'r') as fin:

+            for line in fin:

+                folder, frame_num, _ = line.split(' ')

+                self.keys.extend([f'{folder}/{i:08d}' for i in range(int(frame_num))])

+

+        # remove the video clips used in validation

+        if opt['val_partition'] == 'REDS4':

+            val_partition = ['000', '011', '015', '020']

+        elif opt['val_partition'] == 'official':

+            val_partition = [f'{v:03d}' for v in range(240, 270)]

+        else:

+            raise ValueError(f'Wrong validation partition {opt["val_partition"]}.'

+                             f"Supported ones are ['official', 'REDS4'].")

+        if opt['test_mode']:

+            self.keys = [v for v in self.keys if v.split('/')[0] in val_partition]

+        else:

+            self.keys = [v for v in self.keys if v.split('/')[0] not in val_partition]

+

+        # file client (io backend)

+        self.file_client = None

+        self.io_backend_opt = opt['io_backend']

+        self.is_lmdb = False

+        if self.io_backend_opt['type'] == 'lmdb':

+            self.is_lmdb = True

+            if hasattr(self, 'flow_root') and self.flow_root is not None:

+                self.io_backend_opt['db_paths'] = [self.lq_root, self.gt_root, self.flow_root]

+                self.io_backend_opt['client_keys'] = ['lq', 'gt', 'flow']

+            else:

+                self.io_backend_opt['db_paths'] = [self.lq_root, self.gt_root]

+                self.io_backend_opt['client_keys'] = ['lq', 'gt']

+

+        # temporal augmentation configs

+        self.interval_list = opt.get('interval_list', [1])

+        self.random_reverse = opt.get('random_reverse', False)

+        interval_str = ','.join(str(x) for x in self.interval_list)

+        logger = get_root_logger()

+        logger.info(f'Temporal augmentation interval list: [{interval_str}]; '

+                    f'random reverse is {self.random_reverse}.')

+

+    def __getitem__(self, index):

+        if self.file_client is None:

+            self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)

+

+        scale = self.opt['scale']

+        gt_size = self.opt['gt_size']

+        key = self.keys[index]

+        clip_name, frame_name = key.split('/')  # key example: 000/00000000

+

+        # determine the neighboring frames

+        interval = random.choice(self.interval_list)

+

+        # ensure not exceeding the borders

+        start_frame_idx = int(frame_name)

+        if start_frame_idx > 100 - self.num_frame * interval:

+            start_frame_idx = random.randint(0, 100 - self.num_frame * interval)

+        end_frame_idx = start_frame_idx + self.num_frame * interval

+

+        neighbor_list = list(range(start_frame_idx, end_frame_idx, interval))

+

+        # random reverse

+        if self.random_reverse and random.random() < 0.5:

+            neighbor_list.reverse()

+

+        # get the neighboring LQ and GT frames

+        img_lqs = []

+        img_gts = []

+        for neighbor in neighbor_list:

+            if self.is_lmdb:

+                img_lq_path = f'{clip_name}/{neighbor:08d}'

+                img_gt_path = f'{clip_name}/{neighbor:08d}'

+            else:

+                img_lq_path = self.lq_root / clip_name / f'{neighbor:08d}.png'

+                img_gt_path = self.gt_root / clip_name / f'{neighbor:08d}.png'

+

+            # get LQ

+            img_bytes = self.file_client.get(img_lq_path, 'lq')

+            img_lq = imfrombytes(img_bytes, float32=True)

+            img_lqs.append(img_lq)

+

+            # get GT

+            img_bytes = self.file_client.get(img_gt_path, 'gt')

+            img_gt = imfrombytes(img_bytes, float32=True)

+            img_gts.append(img_gt)

+

+        # randomly crop

+        img_gts, img_lqs = paired_random_crop(img_gts, img_lqs, gt_size, scale, img_gt_path)

+

+        # augmentation - flip, rotate

+        img_lqs.extend(img_gts)

+        img_results = augment(img_lqs, self.opt['use_hflip'], self.opt['use_rot'])

+

+        img_results = img2tensor(img_results)

+        img_gts = torch.stack(img_results[len(img_lqs) // 2:], dim=0)

+        img_lqs = torch.stack(img_results[:len(img_lqs) // 2], dim=0)

+

+        # img_lqs: (t, c, h, w)

+        # img_gts: (t, c, h, w)

+        # key: str

+        return {'lq': img_lqs, 'gt': img_gts, 'key': key}

+

+    def __len__(self):

+        return len(self.keys)