def __getitem__(self, index): folder = self.data_info['folder'][index] idx, max_idx = self.data_info['idx'][index].split('/') idx, max_idx = int(idx), int(max_idx) border = self.data_info['border'][index] lq_path = self.data_info['lq_path'][index] select_idx = util.generate_frame_indices(idx, max_idx, self.opt['num_frame'], padding=self.opt['padding']) if self.cache_data: if self.opt['use_duf_downsampling']: # read imgs_gt to generate low-resolution frames imgs_lq = self.imgs_gt[folder].index_select( 0, torch.LongTensor(select_idx)) imgs_lq = duf_downsample(imgs_lq, kernel_size=13, scale=self.opt['scale']) else: imgs_lq = self.imgs_lq[folder].index_select( 0, torch.LongTensor(select_idx)) img_gt = self.imgs_gt[folder][idx] else: if self.opt['use_duf_downsampling']: img_paths_lq = [self.imgs_gt[folder][i] for i in select_idx] # read imgs_gt to generate low-resolution frames imgs_lq = util.read_img_seq(img_paths_lq, require_mod_crop=True, scale=self.opt['scale']) imgs_lq = duf_downsample(imgs_lq, kernel_size=13, scale=self.opt['scale']) else: img_paths_lq = [self.imgs_lq[folder][i] for i in select_idx] imgs_lq = util.read_img_seq(img_paths_lq) img_gt = util.read_img_seq([self.imgs_gt[folder][idx]], require_mod_crop=True, scale=self.opt['scale']) img_gt.squeeze_(0) return { 'lq': imgs_lq, # (t, c, h, w) 'gt': img_gt, # (c, h, w) 'folder': folder, # folder name 'idx': self.data_info['idx'][index], # e.g., 0/99 'border': border, # 1 for border, 0 for non-border 'lq_path': lq_path # center frame }
def __getitem__(self, index): folder = self.data_info['folder'][index] idx, max_idx = self.data_info['idx'][index].split('/') idx, max_idx = int(idx), int(max_idx) border = self.data_info['border'][index] lq_path = self.data_info['lq_path'][index] select_idx = util.generate_frame_indices(idx, max_idx, self.opt['num_frame'], padding=self.opt['padding']) if self.cache_data: if self.opt['use_duf_downsampling']: imgs_lq = self.imgs_gt[folder].index_select( 0, torch.LongTensor(select_idx)) imgs_lq = duf_downsample(imgs_lq, kernel_size=13, scale=self.opt['scale']) else: imgs_lq = self.imgs_lq[folder].index_select( 0, torch.LongTensor(select_idx)) img_gt = self.imgs_gt[folder][idx] else: if self.opt['use_duf_downsampling']: img_paths_lq = [self.imgs_gt[folder][i] for i in select_idx] imgs_lq = util.read_img_seq(img_paths_lq, require_mod_crop=True, scale=self.opt['scale']) imgs_lq = duf_downsample(imgs_lq, kernel_size=13, scale=self.opt['scale']) else: img_paths_lq = [self.imgs_lq[folder][i] for i in select_idx] imgs_lq = util.read_img_seq(img_paths_lq) img_gt = util.read_img_seq([self.imgs_gt[folder][idx]], require_mod_crop=True, scale=self.opt['scale']) return { 'lq': imgs_lq, 'gt': img_gt, 'folder': folder, 'idx': self.data_info['idx'][index], 'border': border, 'lq_path': lq_path }
def read_img_worker(path, key, compress_level, lr=False, use_wt=False, use_inter_wt=False, filters=None): """Read image worker. Args: path (str): Image path. key (str): Image key. compress_level (int): Compress level when encoding images. Returns: str: Image key. byte: Image byte. tuple[int]: Image shape. """ img = mmcv.imread(path, flag='unchanged') # If lr is True, then downsample from groundtruth images using duf_downsample if lr: img = duf_downsample(torch.from_numpy(img).permute(2,0,1).unsqueeze(0), scale=4).squeeze().permute(1,2,0).numpy() elif use_wt: img = torch.from_numpy(img / 255.0).float() if img.ndim == 2: print(img.shape, 'weird shape of image c=1') sys.exit(1) img = wt(img.unsqueeze(0).unsqueeze(0).to('cuda:0'), filters, levels=3)[:, :, :64, :64].squeeze().cpu().numpy() elif img.ndim == 3: img = wt(img.permute(2,0,1).unsqueeze(0).to('cuda:0'), filters, levels=3)[:, :, :64, :64].squeeze().permute(1,2,0).cpu().numpy() elif use_inter_wt: img = mmcv.image.imresize(img, (128, 128), interpolation='lanczos', backend='cv2') img = torch.from_numpy(img / 255.0).float() if img.ndim == 2: img = wt(img.unsqueeze(0).unsqueeze(0).to('cuda:0'), filters, levels=2)[:, :, :64, :64].squeeze().cpu().numpy() elif img.ndim == 3: img = wt(img.permute(2,0,1).unsqueeze(0).to('cuda:0'), filters, levels=2)[:, :, :64, :64].squeeze().permute(1,2,0).cpu().numpy() if img.ndim == 2: h, w = img.shape c = 1 else: h, w, c = img.shape if not (use_wt or use_inter_wt): _, img_byte = cv2.imencode('.png', img, [cv2.IMWRITE_PNG_COMPRESSION, compress_level]) # If writing WT patch in bytes, use np tobytes() function else: img_byte = img.tobytes() return (key, img_byte, (h, w, c))
def worker(path, opt): """Worker for each process. Args: path (str): Image path. opt (dict): Configuration dict. It contains: crop_size (int): Crop size. save_folder (str): Path to save folder. compression_level (int): for cv2.IMWRITE_PNG_COMPRESSION. Returns: process_info (str): Process information displayed in progress bar. """ scale = opt['scale'] img_name, extension = osp.splitext(osp.basename(path)) extension = '.png' # remove the x2, x3, x4 and x8 in the filename for DIV2K img = cv2.imread(path, cv2.IMREAD_UNCHANGED) if img.ndim == 2: h, w = img.shape elif img.ndim == 3: h, w, c = img.shape else: raise ValueError(f'Image ndim should be 2 or 3, but got {img.ndim}') img = ToTensor()(img) downsampled_img = duf_downsample(img.unsqueeze(0), scale=scale).squeeze(0) downsampled_img *= 255.0 downsampled_img = downsampled_img.numpy().transpose(1,2,0) cv2.imwrite( osp.join(opt['save_folder'], f'{img_name}{extension}'), downsampled_img, [cv2.IMWRITE_PNG_COMPRESSION, opt['compression_level']]) process_info = f'Processing {img_name} ...' return process_info