Updates (XPixelGroup#300)

* update resume pretrained paths

* test_scripts

* add dist util

* use os

* add matlab functions

* scandir and bgr2rgb replace

* cv2.flip

* replace imwrite

* update file client

* update utils.util

* update utils.download

* use relative import

* add img_util

* update strict_load

* updat train.py

* update test.py

* add flow util

* update requirements.txt

* fix bugs

* fix bugs

Author: xinntao

basicsr/data/__init__.py

@@ -1,23 +1,22 @@
 import importlib
-import mmcv
 import numpy as np
 import random
 import torch
 import torch.utils.data
 from functools import partial
-from mmcv.runner import get_dist_info
 from os import path as osp
 
 from basicsr.data.prefetch_dataloader import PrefetchDataLoader
-from basicsr.utils import get_root_logger
+from basicsr.utils import get_root_logger, scandir
+from basicsr.utils.dist_util import get_dist_info
 
 __all__ = ['create_dataset', 'create_dataloader']
 
 # automatically scan and import dataset modules
 # scan all the files under the data folder with '_dataset' in file names
 data_folder = osp.dirname(osp.abspath(__file__))
 dataset_filenames = [
-    osp.splitext(osp.basename(v))[0] for v in mmcv.scandir(data_folder)
+    osp.splitext(osp.basename(v))[0] for v in scandir(data_folder)
     if v.endswith('_dataset.py')
 ]
 # import all the dataset modules

basicsr/data/ffhq_dataset.py

@@ -1,11 +1,9 @@
-import mmcv
-import numpy as np
 from os import path as osp
 from torch.utils import data as data
 from torchvision.transforms.functional import normalize
 
-from basicsr.data.transforms import augment, totensor
-from basicsr.utils import FileClient
+from basicsr.data.transforms import augment
+from basicsr.utils import FileClient, imfrombytes, img2tensor
 
 
 class FFHQDataset(data.Dataset):
@@ -53,12 +51,12 @@ def __getitem__(self, index):
         # load gt image
         gt_path = self.paths[index]
         img_bytes = self.file_client.get(gt_path)
-        img_gt = mmcv.imfrombytes(img_bytes).astype(np.float32) / 255.
+        img_gt = imfrombytes(img_bytes, float32=True)
 
         # random horizontal flip
         img_gt = augment(img_gt, hflip=self.opt['use_hflip'], rotation=False)
         # BGR to RGB, HWC to CHW, numpy to tensor
-        img_gt = totensor(img_gt, bgr2rgb=True, float32=True)
+        img_gt = img2tensor(img_gt, bgr2rgb=True, float32=True)
         # normalize
         normalize(img_gt, self.mean, self.std, inplace=True)
         return {'gt': img_gt, 'gt_path': gt_path}

basicsr/data/paired_image_dataset.py

@@ -1,12 +1,10 @@
-import mmcv
-import numpy as np
 from torch.utils import data as data
 
-from basicsr.data.transforms import augment, paired_random_crop, totensor
+from basicsr.data.transforms import augment, paired_random_crop
 from basicsr.data.util import (paired_paths_from_folder,
                                paired_paths_from_lmdb,
                                paired_paths_from_meta_info_file)
-from basicsr.utils import FileClient
+from basicsr.utils import FileClient, imfrombytes, img2tensor
 
 
 class PairedImageDataset(data.Dataset):
@@ -79,10 +77,10 @@ def __getitem__(self, index):
         # image range: [0, 1], float32.
         gt_path = self.paths[index]['gt_path']
         img_bytes = self.file_client.get(gt_path, 'gt')
-        img_gt = mmcv.imfrombytes(img_bytes).astype(np.float32) / 255.
+        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 = mmcv.imfrombytes(img_bytes).astype(np.float32) / 255.
+        img_lq = imfrombytes(img_bytes, float32=True)
 
         # augmentation for training
         if self.opt['phase'] == 'train':
@@ -96,7 +94,9 @@ def __getitem__(self, index):
 
         # TODO: color space transform
         # BGR to RGB, HWC to CHW, numpy to tensor
-        img_gt, img_lq = totensor([img_gt, img_lq], bgr2rgb=True, float32=True)
+        img_gt, img_lq = img2tensor([img_gt, img_lq],
+                                    bgr2rgb=True,
+                                    float32=True)
 
         return {
             'lq': img_lq,

basicsr/data/reds_dataset.py

@@ -1,12 +1,12 @@
-import mmcv
 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, totensor
-from basicsr.utils import FileClient, get_root_logger
+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
 
 
 class REDSDataset(data.Dataset):
@@ -144,7 +144,7 @@ def __getitem__(self, index):
         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 = mmcv.imfrombytes(img_bytes).astype(np.float32) / 255.
+        img_gt = imfrombytes(img_bytes, float32=True)
 
         # get the neighboring LQ frames
         img_lqs = []
@@ -154,7 +154,7 @@ def __getitem__(self, index):
             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 = mmcv.imfrombytes(img_bytes).astype(np.float32) / 255.
+            img_lq = imfrombytes(img_bytes, float32=True)
             img_lqs.append(img_lq)
 
         # get flows
@@ -168,10 +168,11 @@ def __getitem__(self, index):
                     flow_path = (
                         self.flow_root / clip_name / f'{frame_name}_p{i}.png')
                 img_bytes = self.file_client.get(flow_path, 'flow')
-                cat_flow = mmcv.imfrombytes(
-                    img_bytes, flag='grayscale')  # uint8, [0, 255]
+                cat_flow = imfrombytes(
+                    img_bytes, flag='grayscale',
+                    float32=False)  # uint8, [0, 255]
                 dx, dy = np.split(cat_flow, 2, axis=0)
-                flow = mmcv.video.dequantize_flow(
+                flow = dequantize_flow(
                     dx, dy, max_val=20,
                     denorm=False)  # we use max_val 20 here.
                 img_flows.append(flow)
@@ -183,9 +184,11 @@ def __getitem__(self, index):
                     flow_path = (
                         self.flow_root / clip_name / f'{frame_name}_n{i}.png')
                 img_bytes = self.file_client.get(flow_path, 'flow')
-                cat_flow = mmcv.imfrombytes(img_bytes, flag='grayscale')
+                cat_flow = imfrombytes(
+                    img_bytes, flag='grayscale',
+                    float32=False)  # uint8, [0, 255]
                 dx, dy = np.split(cat_flow, 2, axis=0)
-                flow = mmcv.video.dequantize_flow(
+                flow = dequantize_flow(
                     dx, dy, max_val=20,
                     denorm=False)  # we use max_val 20 here.
                 img_flows.append(flow)
@@ -210,12 +213,12 @@ def __getitem__(self, index):
             img_results = augment(img_lqs, self.opt['use_flip'],
                                   self.opt['use_rot'])
 
-        img_results = totensor(img_results)
+        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 = totensor(img_flows)
+            img_flows = img2tensor(img_flows)
             # add the zero center flow
             img_flows.insert(self.num_half_frames,
                              torch.zeros_like(img_flows[0]))

basicsr/data/single_image_dataset.py

@@ -1,11 +1,8 @@
-import mmcv
-import numpy as np
 from os import path as osp
 from torch.utils import data as data
 from torchvision.transforms.functional import normalize
 
-from basicsr.data.transforms import totensor
-from basicsr.utils import FileClient
+from basicsr.utils import FileClient, imfrombytes, img2tensor, scandir
 
 
 class SingleImageDataset(data.Dataset):
@@ -40,10 +37,7 @@ def __init__(self, opt):
                              line.split(' ')[0]) for line in fin
                 ]
         else:
-            self.paths = [
-                osp.join(self.lq_folder, v)
-                for v in mmcv.scandir(self.lq_folder)
-            ]
+            self.paths = sorted(list(scandir(self.lq_folder, full_path=True)))
 
     def __getitem__(self, index):
         if self.file_client is None:
@@ -53,11 +47,11 @@ def __getitem__(self, index):
         # load lq image
         lq_path = self.paths[index]
         img_bytes = self.file_client.get(lq_path)
-        img_lq = mmcv.imfrombytes(img_bytes).astype(np.float32) / 255.
+        img_lq = imfrombytes(img_bytes, float32=True)
 
         # TODO: color space transform
         # BGR to RGB, HWC to CHW, numpy to tensor
-        img_lq = totensor(img_lq, bgr2rgb=True, float32=True)
+        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)

basicsr/data/transforms.py

@@ -1,6 +1,5 @@
-import mmcv
+import cv2
 import random
-import torch
 
 
 def mod_crop(img, scale):
@@ -110,20 +109,20 @@ def augment(imgs, hflip=True, rotation=True, flows=None):
     rot90 = rotation and random.random() < 0.5
 
     def _augment(img):
-        if hflip:
-            mmcv.imflip_(img, 'horizontal')
-        if vflip:
-            mmcv.imflip_(img, 'vertical')
+        if hflip:  # horizontal
+            cv2.flip(img, 1, img)
+        if vflip:  # vertical
+            cv2.flip(img, 0, img)
         if rot90:
             img = img.transpose(1, 0, 2)
         return img
 
     def _augment_flow(flow):
-        if hflip:
-            mmcv.imflip_(flow, 'horizontal')
+        if hflip:  # horizontal
+            cv2.flip(flow, 1, flow)
             flow[:, :, 0] *= -1
-        if vflip:
-            mmcv.imflip_(flow, 'vertical')
+        if vflip:  # vertical
+            cv2.flip(flow, 0, flow)
             flow[:, :, 1] *= -1
         if rot90:
             flow = flow.transpose(1, 0, 2)
@@ -145,30 +144,3 @@ def _augment_flow(flow):
         return imgs, flows
     else:
         return imgs
-
-
-def totensor(imgs, bgr2rgb=True, float32=True):
-    """Numpy array to tensor.
-
-    Args:
-        imgs (list[ndarray] | ndarray): Input images.
-        bgr2rgb (bool): Whether to change bgr to rgb.
-        float32 (bool): Whether to change to float32.
-
-    Returns:
-        list[tensor] | tensor: Tensor images. If returned results only have
-            one element, just return tensor.
-    """
-
-    def _totensor(img, bgr2rgb, float32):
-        if img.shape[2] == 3 and bgr2rgb:
-            img = mmcv.bgr2rgb(img)
-        img = torch.from_numpy(img.transpose(2, 0, 1))
-        if float32:
-            img = img.float()
-        return img
-
-    if isinstance(imgs, list):
-        return [_totensor(img, bgr2rgb, float32) for img in imgs]
-    else:
-        return _totensor(imgs, bgr2rgb, float32)

basicsr/data/util.py

@@ -1,10 +1,11 @@
-import mmcv
+import cv2
 import numpy as np
 import torch
 from os import path as osp
 from torch.nn import functional as F
 
-from basicsr.data.transforms import mod_crop, totensor
+from basicsr.data.transforms import mod_crop
+from basicsr.utils import img2tensor, scandir
 
 
 def read_img_seq(path, require_mod_crop=False, scale=1):
@@ -22,11 +23,11 @@ def read_img_seq(path, require_mod_crop=False, scale=1):
     if isinstance(path, list):
         img_paths = path
     else:
-        img_paths = sorted([osp.join(path, v) for v in mmcv.scandir(path)])
-    imgs = [mmcv.imread(v).astype(np.float32) / 255. for v in img_paths]
+        img_paths = sorted(list(scandir(path, full_path=True)))
+    imgs = [cv2.imread(v).astype(np.float32) / 255. for v in img_paths]
     if require_mod_crop:
         imgs = [mod_crop(img, scale) for img in imgs]
-    imgs = totensor(imgs, bgr2rgb=True, float32=True)
+    imgs = img2tensor(imgs, bgr2rgb=True, float32=True)
     imgs = torch.stack(imgs, dim=0)
     return imgs
 
@@ -227,8 +228,8 @@ def paired_paths_from_folder(folders, keys, filename_tmpl):
     input_folder, gt_folder = folders
     input_key, gt_key = keys
 
-    input_paths = list(mmcv.scandir(input_folder))
-    gt_paths = list(mmcv.scandir(gt_folder))
+    input_paths = list(scandir(input_folder))
+    gt_paths = list(scandir(gt_folder))
     assert len(input_paths) == len(gt_paths), (
         f'{input_key} and {gt_key} datasets have different number of images: '
         f'{len(input_paths)}, {len(gt_paths)}.')
@@ -256,7 +257,7 @@ def paths_from_folder(folder):
         list[str]: Returned path list.
     """
 
-    paths = list(mmcv.scandir(folder))
+    paths = list(scandir(folder))
     paths = [osp.join(folder, path) for path in paths]
     return paths
 

basicsr/data/video_test_dataset.py

@@ -1,12 +1,11 @@
 import glob
-import mmcv
 import torch
 from os import path as osp
 from torch.utils import data as data
 
 from basicsr.data import util as util
 from basicsr.data.util import duf_downsample
-from basicsr.utils import get_root_logger
+from basicsr.utils import get_root_logger, scandir
 
 
 class VideoTestDataset(data.Dataset):
@@ -81,14 +80,10 @@ def __init__(self, opt):
                                                   subfolders_gt):
                 # get frame list for lq and gt
                 subfolder_name = osp.basename(subfolder_lq)
-                img_paths_lq = sorted([
-                    osp.join(subfolder_lq, v)
-                    for v in mmcv.scandir(subfolder_lq)
-                ])
-                img_paths_gt = sorted([
-                    osp.join(subfolder_gt, v)
-                    for v in mmcv.scandir(subfolder_gt)
-                ])
+                img_paths_lq = sorted(
+                    list(scandir(subfolder_lq, full_path=True)))
+                img_paths_gt = sorted(
+                    list(scandir(subfolder_gt, full_path=True)))
 
                 max_idx = len(img_paths_lq)
                 assert max_idx == len(img_paths_gt), (