text: import mmcv
  import numpy as np
  import random
  import torch
  from PIL import Image
  from pathlib import Path
  from torch.utils import data as data
  
  from basicsr.data.transforms import augment, paired_random_crop, totensor
  from basicsr.utils import FileClient, get_root_logger
  
  
  class Vimeo90KDataset(data.Dataset):
      """Vimeo90K dataset for training.
  
      The keys are generated from a meta info txt file.
      basicsr/data/meta_info/meta_info_Vimeo90K_train_GT.txt
  
      Each line contains:
      1. clip name; 2. frame number; 3. image shape, seperated by a white space.
      Examples:
          00001/0001 7 (256,448,3)
          00001/0002 7 (256,448,3)
  
      Key examples: "00001/0001"
      GT (gt): Ground-Truth;
      LQ (lq): Low-Quality, e.g., low-resolution/blurry/noisy/compressed frames.
  
      The neighboring frame list for different num_frame:
      num_frame | frame list
               1 | 4
               3 | 3,4,5
               5 | 2,3,4,5,6
               7 | 1,2,3,4,5,6,7
  
      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.
              meta_info_file (str): Path for meta information file.
              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.
              random_reverse (bool): Random reverse input frames.
              use_flip (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(Vimeo90KDataset, self).__init__()
          self.opt = opt
          self.gt_root, self.lq_root = Path(opt['dataroot_gt']), Path(
              opt['dataroot_lq'])
  
          with open(opt['meta_info_file'], 'r') as fin:
              self.keys = [line.split(' ')[0] for line in fin]
  
          # 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
              self.io_backend_opt['db_paths'] = [self.lq_root, self.gt_root]
              self.io_backend_opt['client_keys'] = ['lq', 'gt']
  
          # indices of input images
          self.neighbor_list = [
              i + (9 - opt['num_frame']) // 2 for i in range(opt['num_frame'])
          ]
  
          # temporal augmentation configs
          self.random_reverse = False
          logger = get_root_logger()
          logger.info(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)
  
          # random reverse
          if self.random_reverse and random.random() < 0.5:
              self.neighbor_list.reverse()
  
          scale = self.opt['scale']
          gt_size = self.opt['gt_size']
          key = self.keys[index]
          clip, seq = key.split('/')  # key example: 00001/0001
  
          # get the GT frame (im4.png)
          if self.is_lmdb:
              img_gt_path = f'{key}/im4'
          else:
              img_gt_path = self.gt_root / clip / seq / 'im4.png'
          img_bytes = self.file_client.get(img_gt_path, 'gt')
          img_gt = mmcv.imfrombytes(img_bytes).astype(np.float32) / 255.
          ### get 160
          img_160_path = self.lq_root / clip / seq / 'im4_160.png'
          img_bytes_160 = self.file_client.get(img_160_path, 'gt')
          img_160 = mmcv.imfrombytes(img_bytes_160).astype(np.float32) / 255.
          # get the neighboring LQ frames
          img_gt_160 = []
          img_gt_160.append(img_gt)
          img_gt_160.append(img_160)
  #         ###visualization
  #         path='/home/wei/exp/EDVR/visualization'
  #         number = 1
  #         ###visualization
          img_lqs = []
          for neighbor in self.neighbor_list:
              if self.is_lmdb:
                  img_lq_path = f'{clip}/{seq}/im{neighbor}'
              else:
                  img_lq_path = self.lq_root / clip / seq / f'im{neighbor}.png'
              img_bytes = self.file_client.get(img_lq_path, 'lq')
              img_lq = mmcv.imfrombytes(img_bytes).astype(np.float32) / 255.
            #  img_con_3d = np.vstack((img_lq, img_3d))
  #             ##visualization
  #             number +=1
              
  #             visual_lq = img_con_3d[:,:,:3]
  #             visual_lq = Image.fromarray((visual_lq).astype(np.uint8)).convert("RGB")
  #             visual_lq.save(path+'/'+str(number)+'_lq.png')
  #             visual_3d = img_con_3d[:,:,3:]
  #             visual_3d = Image.fromarray((visual_3d).astype(np.uint8)).convert("RGB")
  #             visual_3d.save(path+'/'+str(number)+'_3d.png')
  #             ##visualization
              img_lqs.append(img_lq)
              
              
              
          # randomly crop
          img_gt, img_lqs = paired_random_crop(img_gt_160, img_lqs, gt_size, scale,
                                               img_gt_path)
          img_160_input = img_gt[1]
          img_gt = img_gt[0]
          # augmentation - flip, rotate
          img_lqs.append(img_160_input)
          img_lqs.append(img_gt)
          img_results = augment(img_lqs, self.opt['use_flip'],
                                self.opt['use_rot'])
  
          img_results = totensor(img_results)
          hr_3d = img_results[-2]
          img_lqs = torch.stack(img_results[0:-2], dim=0)
          img_gt = img_results[-1]
  
          # img_lqs: (t, c, h, w)
          # img_gt: (c, h, w)
          # key: str
          return {'lq': img_lqs, 'gt': img_gt, 'hr_3d':hr_3d, 'key': key}
  
      def __len__(self):
          return len(self.keys)
  