Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import random | |
| import torch | |
| from pathlib import Path | |
| from torch.utils import data as data | |
| from basicsr.data.transforms import augment, paired_random_crop | |
| from basicsr.utils import FileClient, get_root_logger, imfrombytes, img2tensor | |
| from basicsr.utils.flow_util import dequantize_flow | |
| from basicsr.utils.registry import DATASET_REGISTRY | |
| 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) | |
| 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) | |