ICCV 2023

Former-commit-id: af824155fb8ed74383da1f2edf27b09dfab3e6fe

basicsr/data/data_util.py

@@ -162,7 +162,7 @@ def paired_paths_from_meta_info_file(folders, keys, meta_info_file, filename_tmp
     return paths
 
 
-def paired_paths_from_folder(folders, keys, filename_tmpl, task):
+def paired_paths_from_folder(folders, keys, filename_tmpl):
     """Generate paired paths from folders.
 
     Args:
@@ -190,10 +190,7 @@ def paired_paths_from_folder(folders, keys, filename_tmpl, task):
     paths = []
     for gt_path in gt_paths:
         basename, ext = osp.splitext(osp.basename(gt_path))
-        if task == "CAR":
-            input_name = f'{filename_tmpl.format(basename)}.jpg'
-        else:
-            input_name = f'{filename_tmpl.format(basename)}{ext}'
+        input_name = f'{filename_tmpl.format(basename)}{ext}'
         input_path = osp.join(input_folder, input_name)
         assert input_name in input_paths, f'{input_name} is not in {input_key}_paths.'
         gt_path = osp.join(gt_folder, gt_path)

basicsr/data/paired_image_dataset.py

@@ -46,7 +46,6 @@ class PairedImageDataset(data.Dataset):
         self.file_client = None
         self.io_backend_opt = opt['io_backend']
         self.mean = opt['mean'] if 'mean' in opt else None
-        self.task = opt['task'] if 'task' in opt else None
         self.std = opt['std'] if 'std' in opt else None
 
         self.gt_folder, self.lq_folder = opt['dataroot_gt'], opt['dataroot_lq']
@@ -63,7 +62,7 @@ class PairedImageDataset(data.Dataset):
             self.paths = paired_paths_from_meta_info_file([self.lq_folder, self.gt_folder], ['lq', 'gt'],
                                                           self.opt['meta_info_file'], self.filename_tmpl)
         else:
-            self.paths = paired_paths_from_folder([self.lq_folder, self.gt_folder], ['lq', 'gt'], self.filename_tmpl, self.task)
+            self.paths = paired_paths_from_folder([self.lq_folder, self.gt_folder], ['lq', 'gt'], self.filename_tmpl)
 
     def __getitem__(self, index):
         if self.file_client is None:
@@ -72,35 +71,14 @@ class PairedImageDataset(data.Dataset):
         scale = self.opt['scale']
 
         # Load gt and lq images. Dimension order: HWC; channel order: BGR;
-
-        if self.task == 'CAR':
-            # image range: [0, 255], int., H W 1
-            gt_path = self.paths[index]['gt_path']
-            img_bytes = self.file_client.get(gt_path, 'gt')
-            img_gt = imfrombytes(img_bytes, flag='grayscale', float32=False)
-            lq_path = self.paths[index]['lq_path']
-            img_bytes = self.file_client.get(lq_path, 'lq')
-            img_lq = imfrombytes(img_bytes, flag='grayscale', float32=False)
-            img_gt = np.expand_dims(img_gt, axis=2).astype(np.float32) / 255.
-            img_lq = np.expand_dims(img_lq, axis=2).astype(np.float32) / 255.
-
-        elif self.task == 'Color-DN':
-            gt_path = self.paths[index]['gt_path']
-            lq_path = gt_path
-            img_bytes = self.file_client.get(gt_path, 'gt')
-            img_gt = imfrombytes(img_bytes, float32=True)
-            if self.opt['phase'] != 'train':
-                np.random.seed(seed=0)
-            img_lq = img_gt + np.random.normal(0, self.noise/255., img_gt.shape)
-
-        else:
-            # image range: [0, 1], float32., H W 3
-            gt_path = self.paths[index]['gt_path']
-            img_bytes = self.file_client.get(gt_path, 'gt')
-            img_gt = imfrombytes(img_bytes, float32=True)
-            lq_path = self.paths[index]['lq_path']
-            img_bytes = self.file_client.get(lq_path, 'lq')
-            img_lq = imfrombytes(img_bytes, float32=True)
+        
+        # image range: [0, 1], float32., H W 3
+        gt_path = self.paths[index]['gt_path']
+        img_bytes = self.file_client.get(gt_path, 'gt')
+        img_gt = imfrombytes(img_bytes, float32=True)
+        lq_path = self.paths[index]['lq_path']
+        img_bytes = self.file_client.get(lq_path, 'lq')
+        img_lq = imfrombytes(img_bytes, float32=True)
 
         # augmentation for training
         if self.opt['phase'] == 'train':

basicsr/data/single_image_dataset.py

@@ -0,0 +1,69 @@
+from os import path as osp
+from torch.utils import data as data
+from torchvision.transforms.functional import normalize
+
+from basicsr.data.data_util import paths_from_lmdb
+from basicsr.utils import FileClient, imfrombytes, img2tensor, scandir
+from basicsr.utils.matlab_functions import rgb2ycbcr
+from basicsr.utils.registry import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register()
+class SingleImageDataset(data.Dataset):
+    """Read only lq images in the test phase.
+
+    Read LQ (Low Quality, e.g. LR (Low Resolution), blurry, noisy, etc).
+
+    There are two modes:
+    1. 'meta_info_file': Use meta information file to generate paths.
+    2. 'folder': Scan folders to generate paths.
+
+    Args:
+        opt (dict): Config for train datasets. It contains the following keys:
+            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.
+    """
+
+    def __init__(self, opt):
+        super(SingleImageDataset, self).__init__()
+        self.opt = opt
+        # file client (io backend)
+        self.file_client = None
+        self.io_backend_opt = opt['io_backend']
+        self.mean = opt['mean'] if 'mean' in opt else None
+        self.std = opt['std'] if 'std' in opt else None
+        self.lq_folder = opt['dataroot_lq']
+
+        if self.io_backend_opt['type'] == 'lmdb':
+            self.io_backend_opt['db_paths'] = [self.lq_folder]
+            self.io_backend_opt['client_keys'] = ['lq']
+            self.paths = paths_from_lmdb(self.lq_folder)
+        elif 'meta_info_file' in self.opt:
+            with open(self.opt['meta_info_file'], 'r') as fin:
+                self.paths = [osp.join(self.lq_folder, line.rstrip().split(' ')[0]) for line in fin]
+        else:
+            self.paths = sorted(list(scandir(self.lq_folder, full_path=True)))
+
+    def __getitem__(self, index):
+        if self.file_client is None:
+            self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
+
+        # load lq image
+        lq_path = self.paths[index]
+        img_bytes = self.file_client.get(lq_path, 'lq')
+        img_lq = imfrombytes(img_bytes, float32=True)
+
+        # color space transform
+        if 'color' in self.opt and self.opt['color'] == 'y':
+            img_lq = rgb2ycbcr(img_lq, y_only=True)[..., None]
+
+        # BGR to RGB, HWC to CHW, numpy to tensor
+        img_lq = img2tensor(img_lq, bgr2rgb=True, float32=True)
+        # normalize
+        if self.mean is not None or self.std is not None:
+            normalize(img_lq, self.mean, self.std, inplace=True)
+        return {'lq': img_lq, 'lq_path': lq_path}
+
+    def __len__(self):
+        return len(self.paths)

basicsr/models/dat_model.py

@@ -0,0 +1,127 @@
+import torch
+from torch.nn import functional as F
+
+from basicsr.utils.registry import MODEL_REGISTRY
+from basicsr.models.sr_model import SRModel
+
+
+@MODEL_REGISTRY.register()
+class DATModle(SRModel):
+
+    def test(self):
+        self.use_chop = self.opt['val']['use_chop'] if 'use_chop' in self.opt['val'] else False
+        if not self.use_chop:
+            if hasattr(self, 'net_g_ema'):
+                self.net_g_ema.eval()
+                with torch.no_grad():
+                    self.output = self.net_g_ema(self.lq)
+            else:
+                self.net_g.eval()
+                with torch.no_grad():
+                    self.output = self.net_g(self.lq)
+                self.net_g.train()
+
+        # test by partitioning
+        else:
+            _, C, h, w = self.lq.size()
+            split_token_h = h // 200 + 1  # number of horizontal cut sections
+            split_token_w = w // 200 + 1  # number of vertical cut sections
+
+            patch_size_tmp_h = split_token_h
+            patch_size_tmp_w = split_token_w
+            
+            # padding
+            mod_pad_h, mod_pad_w = 0, 0
+            if h % patch_size_tmp_h != 0:
+                mod_pad_h = patch_size_tmp_h - h % patch_size_tmp_h
+            if w % patch_size_tmp_w != 0:
+                mod_pad_w = patch_size_tmp_w - w % patch_size_tmp_w
+                
+            img = self.lq
+            img = torch.cat([img, torch.flip(img, [2])], 2)[:, :, :h+mod_pad_h, :]
+            img = torch.cat([img, torch.flip(img, [3])], 3)[:, :, :, :w+mod_pad_w]
+
+            _, _, H, W = img.size()
+            split_h = H // split_token_h  # height of each partition
+            split_w = W // split_token_w  # width of each partition
+
+            # overlapping
+            shave_h = 16
+            shave_w = 16
+            scale = self.opt.get('scale', 1)
+            ral = H // split_h
+            row = W // split_w
+            slices = []  # list of partition borders
+            for i in range(ral):
+                for j in range(row):
+                    if i == 0 and i == ral - 1:
+                        top = slice(i * split_h, (i + 1) * split_h)
+                    elif i == 0:
+                        top = slice(i*split_h, (i+1)*split_h+shave_h)
+                    elif i == ral - 1:
+                        top = slice(i*split_h-shave_h, (i+1)*split_h)
+                    else:
+                        top = slice(i*split_h-shave_h, (i+1)*split_h+shave_h)
+                    if j == 0 and j == row - 1:
+                        left = slice(j*split_w, (j+1)*split_w)
+                    elif j == 0:
+                        left = slice(j*split_w, (j+1)*split_w+shave_w)
+                    elif j == row - 1:
+                        left = slice(j*split_w-shave_w, (j+1)*split_w)
+                    else:
+                        left = slice(j*split_w-shave_w, (j+1)*split_w+shave_w)
+                    temp = (top, left)
+                    slices.append(temp)
+            img_chops = []  # list of partitions
+            for temp in slices:
+                top, left = temp
+                img_chops.append(img[..., top, left])
+            if hasattr(self, 'net_g_ema'):
+                self.net_g_ema.eval()
+                with torch.no_grad():
+                    outputs = []
+                    for chop in img_chops:
+                        out = self.net_g_ema(chop)  # image processing of each partition
+                        outputs.append(out)
+                    _img = torch.zeros(1, C, H * scale, W * scale)
+                    # merge
+                    for i in range(ral):
+                        for j in range(row):
+                            top = slice(i * split_h * scale, (i + 1) * split_h * scale)
+                            left = slice(j * split_w * scale, (j + 1) * split_w * scale)
+                            if i == 0:
+                                _top = slice(0, split_h * scale)
+                            else:
+                                _top = slice(shave_h*scale, (shave_h+split_h)*scale)
+                            if j == 0:
+                                _left = slice(0, split_w*scale)
+                            else:
+                                _left = slice(shave_w*scale, (shave_w+split_w)*scale)
+                            _img[..., top, left] = outputs[i * row + j][..., _top, _left]
+                    self.output = _img
+            else:
+                self.net_g.eval()
+                with torch.no_grad():
+                    outputs = []
+                    for chop in img_chops:
+                        out = self.net_g(chop)  # image processing of each partition
+                        outputs.append(out)
+                    _img = torch.zeros(1, C, H * scale, W * scale)
+                    # merge
+                    for i in range(ral):
+                        for j in range(row):
+                            top = slice(i * split_h * scale, (i + 1) * split_h * scale)
+                            left = slice(j * split_w * scale, (j + 1) * split_w * scale)
+                            if i == 0:
+                                _top = slice(0, split_h * scale)
+                            else:
+                                _top = slice(shave_h * scale, (shave_h + split_h) * scale)
+                            if j == 0:
+                                _left = slice(0, split_w * scale)
+                            else:
+                                _left = slice(shave_w * scale, (shave_w + split_w) * scale)
+                            _img[..., top, left] = outputs[i * row + j][..., _top, _left]
+                    self.output = _img
+                self.net_g.train()
+            _, _, h, w = self.output.size()
+            self.output = self.output[:, :, 0:h - mod_pad_h * scale, 0:w - mod_pad_w * scale]

datasets/README.md

@@ -2,6 +2,9 @@ For training and testing, the directory structure is as follows:
 
 ```shell
 |-- datasets
+    # test your own imgs
+    |-- single
+
     # train
     |-- DF2K
         |-- HR

datasets/single/.gitattributes

@@ -0,0 +1 @@
+*.{png,jpg} filter=lfs diff=lfs merge=lfs -text

options/Test/test_DAT_2_x2.yml

@@ -1,6 +1,6 @@
 # general settings
 name: test_DAT_2_x2
-model_type: SRModel
+model_type: DATModle
 scale: 2
 num_gpu: 1
 manual_seed: 10
@@ -80,7 +80,7 @@ path:
 val:
   save_img: True
   suffix: ~  # add suffix to saved images, if None, use exp name
-  use_chop: False
+  use_chop: False  # True to save memory, if img too large
 
   metrics:
     psnr: # metric name, can be arbitrary

options/Test/test_DAT_2_x3.yml

@@ -1,6 +1,6 @@
 # general settings
 name: test_DAT_2_x3
-model_type: SRModel
+model_type: DATModle
 scale: 3
 num_gpu: 1
 manual_seed: 10
@@ -79,7 +79,7 @@ path:
 val:
   save_img: True
   suffix: ~  # add suffix to saved images, if None, use exp name
-  use_chop: False
+  use_chop: False  # True to save memory, if img too large
 
   metrics:
     psnr: # metric name, can be arbitrary

options/Test/test_DAT_2_x4.yml

@@ -1,6 +1,6 @@
 # general settings
 name: test_DAT_2_x4
-model_type: SRModel
+model_type: DATModle
 scale: 4
 num_gpu: 1
 manual_seed: 10
@@ -80,7 +80,7 @@ path:
 val:
   save_img: True
   suffix: ~  # add suffix to saved images, if None, use exp name
-  use_chop: False
+  use_chop: False  # True to save memory, if img too large
 
   metrics:
     psnr: # metric name, can be arbitrary

options/Test/test_DAT_S_x2.yml

@@ -1,6 +1,6 @@
 # general settings
 name: test_DAT_S_x2
-model_type: SRModel
+model_type: DATModle
 scale: 2
 num_gpu: 1
 manual_seed: 10
@@ -80,7 +80,7 @@ path:
 val:
   save_img: True
   suffix: ~  # add suffix to saved images, if None, use exp name
-  use_chop: False
+  use_chop: False  # True to save memory, if img too large
 
   metrics:
     psnr: # metric name, can be arbitrary

options/Test/test_DAT_S_x3.yml

@@ -1,6 +1,6 @@
 # general settings
 name: test_DAT_S_x3
-model_type: SRModel
+model_type: DATModle
 scale: 3
 num_gpu: 1
 manual_seed: 10
@@ -79,7 +79,7 @@ path:
 val:
   save_img: True
   suffix: ~  # add suffix to saved images, if None, use exp name
-  use_chop: False
+  use_chop: False  # True to save memory, if img too large
 
   metrics:
     psnr: # metric name, can be arbitrary

options/Test/test_DAT_S_x4.yml

@@ -1,6 +1,6 @@
 # general settings
 name: test_DAT_S_x4
-model_type: SRModel
+model_type: DATModle
 scale: 4
 num_gpu: 1
 manual_seed: 10
@@ -80,7 +80,7 @@ path:
 val:
   save_img: True
   suffix: ~  # add suffix to saved images, if None, use exp name
-  use_chop: False
+  use_chop: False  # True to save memory, if img too large
 
   metrics:
     psnr: # metric name, can be arbitrary

options/Test/test_DAT_x2.yml

@@ -1,6 +1,6 @@
 # general settings
 name: test_DAT_x2
-model_type: SRModel
+model_type: DATModle
 scale: 2
 num_gpu: 1
 manual_seed: 10
@@ -80,7 +80,7 @@ path:
 val:
   save_img: True
   suffix: ~  # add suffix to saved images, if None, use exp name
-  use_chop: False
+  use_chop: False  # True to save memory, if img too large
 
   metrics:
     psnr: # metric name, can be arbitrary

options/Test/test_DAT_x3.yml

@@ -1,6 +1,6 @@
 # general settings
 name: test_DAT_x3
-model_type: SRModel
+model_type: DATModle
 scale: 3
 num_gpu: 1
 manual_seed: 10
@@ -79,7 +79,7 @@ path:
 val:
   save_img: True
   suffix: ~  # add suffix to saved images, if None, use exp name
-  use_chop: False
+  use_chop: False  # True to save memory, if img too large
 
   metrics:
     psnr: # metric name, can be arbitrary

options/Test/test_DAT_x4.yml

@@ -1,6 +1,6 @@
 # general settings
 name: test_DAT_x4
-model_type: SRModel
+model_type: DATModle
 scale: 4
 num_gpu: 1
 manual_seed: 10
@@ -80,7 +80,7 @@ path:
 val:
   save_img: True
   suffix: ~  # add suffix to saved images, if None, use exp name
-  use_chop: False
+  use_chop: False  # True to save memory, if img too large
 
   metrics:
     psnr: # metric name, can be arbitrary

options/Test/test_single_x2.yml

@@ -0,0 +1,40 @@
+# general settings
+name: test_single_x2
+model_type: DATModle
+scale: 2
+num_gpu: 1
+manual_seed: 10
+
+datasets:
+  test_1:  # the 1st test dataset
+    name: Single
+    type: SingleImageDataset
+    dataroot_lq: datasets/single
+    io_backend:
+      type: disk
+
+
+# network structures
+network_g:
+  type: DAT
+  upscale: 2
+  in_chans: 3
+  img_size: 64
+  img_range: 1.
+  split_size: [8,32]
+  depth: [6,6,6,6,6,6]
+  embed_dim: 180
+  num_heads: [6,6,6,6,6,6]
+  expansion_factor: 4
+  resi_connection: '1conv'
+
+# path
+path:
+  pretrain_network_g: experiments/pretrained_models/DAT/DAT_x2.pth
+  strict_load_g: True
+
+# validation settings
+val:
+  save_img: True
+  suffix: 'x2'  # add suffix to saved images, if None, use exp name
+  use_chop: False  # True to save memory, if img too large

options/Test/test_single_x3.yml

@@ -0,0 +1,40 @@
+# general settings
+name: test_single_x3
+model_type: DATModle
+scale: 3
+num_gpu: 1
+manual_seed: 10
+
+datasets:
+  test_1:  # the 1st test dataset
+    name: Single
+    type: SingleImageDataset
+    dataroot_lq: datasets/single
+    io_backend:
+      type: disk
+
+
+# network structures
+network_g:
+  type: DAT
+  upscale: 3
+  in_chans: 3
+  img_size: 64
+  img_range: 1.
+  split_size: [8,32]
+  depth: [6,6,6,6,6,6]
+  embed_dim: 180
+  num_heads: [6,6,6,6,6,6]
+  expansion_factor: 4
+  resi_connection: '1conv'
+
+# path
+path:
+  pretrain_network_g: experiments/pretrained_models/DAT/DAT_x3.pth
+  strict_load_g: True
+
+# validation settings
+val:
+  save_img: True
+  suffix: 'x3'  # add suffix to saved images, if None, use exp name
+  use_chop: False  # True to save memory, if img too large

options/Test/test_single_x4.yml

@@ -0,0 +1,40 @@
+# general settings
+name: test_single_x4
+model_type: DATModle
+scale: 4
+num_gpu: 1
+manual_seed: 10
+
+datasets:
+  test_1:  # the 1st test dataset
+    name: Single
+    type: SingleImageDataset
+    dataroot_lq: datasets/single
+    io_backend:
+      type: disk
+
+
+# network structures
+network_g:
+  type: DAT
+  upscale: 4
+  in_chans: 3
+  img_size: 64
+  img_range: 1.
+  split_size: [8,32]
+  depth: [6,6,6,6,6,6]
+  embed_dim: 180
+  num_heads: [6,6,6,6,6,6]
+  expansion_factor: 4
+  resi_connection: '1conv'
+
+# path
+path:
+  pretrain_network_g: experiments/pretrained_models/DAT/DAT_x4.pth
+  strict_load_g: True
+
+# validation settings
+val:
+  save_img: True
+  suffix: 'x4'  # add suffix to saved images, if None, use exp name
+  use_chop: False  # True to save memory, if img too large

options/Train/train_DAT_2_x2.yml

@@ -1,6 +1,6 @@
 # general settings
 name: train_DAT_2_x2
-model_type: SRModel
+model_type: DATModle
 scale: 2
 num_gpu: auto
 manual_seed: 10

options/Train/train_DAT_2_x3.yml

@@ -1,6 +1,6 @@
 # general settings
 name: train_DAT_2_x3
-model_type: SRModel
+model_type: DATModle
 scale: 3
 num_gpu: auto
 manual_seed: 10

options/Train/train_DAT_2_x4.yml

@@ -1,6 +1,6 @@
 # general settings
 name: test_DAT_2_x4
-model_type: SRModel
+model_type: DATModle
 scale: 4
 num_gpu: auto
 manual_seed: 10

options/Train/train_DAT_S_x2.yml

@@ -1,6 +1,6 @@
 # general settings
 name: train_DAT_S_x2
-model_type: SRModel
+model_type: DATModle
 scale: 2
 num_gpu: auto
 manual_seed: 10
@@ -24,7 +24,7 @@ datasets:
     # data loader
     use_shuffle: True
     num_worker_per_gpu: 12
-    batch_size_per_gpu: 8
+    batch_size_per_gpu: 2
     dataset_enlarge_ratio: 100
     prefetch_mode: ~
 

options/Train/train_DAT_S_x3.yml

@@ -1,6 +1,6 @@
 # general settings
 name: train_DAT_S_x3
-model_type: SRModel
+model_type: DATModle
 scale: 3
 num_gpu: auto
 manual_seed: 10

options/Train/train_DAT_S_x4.yml

@@ -1,6 +1,6 @@
 # general settings
 name: test_DAT_S_x4
-model_type: SRModel
+model_type: DATModle
 scale: 4
 num_gpu: auto
 manual_seed: 10

options/Train/train_DAT_x2.yml

@@ -1,6 +1,6 @@
 # general settings
 name: train_DAT_x2
-model_type: SRModel
+model_type: DATModle
 scale: 2
 num_gpu: auto
 manual_seed: 10

options/Train/train_DAT_x3.yml

@@ -1,6 +1,6 @@
 # general settings
 name: train_DAT_x3
-model_type: SRModel
+model_type: DATModle
 scale: 3
 num_gpu: auto
 manual_seed: 10

options/Train/train_DAT_x4.yml

@@ -1,6 +1,6 @@
 # general settings
 name: test_DAT_x4
-model_type: SRModel
+model_type: DATModle
 scale: 4
 num_gpu: auto
 manual_seed: 10