first commit

Author: cszhilu1998

README.md

@@ -0,0 +1,95 @@
+# SelfIR (NeurIPS 2022)
+
+PyTorch implementation of [**Self-Supervised Image Restoration with Blurry and Noisy Pairs**](https://arxiv.org/) 
+
+[Paper (arXiv)](https://arxiv.org/)        
+[Video](https://arxiv.org/)
+
+## 1. Framework
+
+<p align="center"><img src="overview.png" width="95%"></p>
+<p align="center">Overview of our proposed SelfIR framework.</p>
+
+- (a) Training phase of SelfIR. Sub-sampled blurry image $\mathit{g}_1(\mathbf{I}_\mathcal{B})$ and noisy image $\mathit{g}_1(\mathbf{I}_\mathcal{N})$ are taken as the inputs. $\mathit{g}_2(\mathbf{I}_\mathcal{N})$ is used for calculating the reconstruction loss $\mathcal{L}_\mathit{rec}$ and regularization loss $\mathcal{L}_\mathit{reg}$, while $\mathit{g}_1(\mathbf{I}_\mathcal{B})$ is taken for calculating auxiliary loss $\mathcal{L}_\mathit{aux}$.
+- (b) Example of neighbor sub-sampler. In each $2\times2$ cell, two pixels are randomly selected for respectively composing the neighboring sub-images.
+- (c) Testing phase of SelfIR. The blurry and noisy images can be directly taken for restoration.
+
+## 2. Preparation and Datasets
+
+- **Prerequisites**
+    - Python 3.x and **PyTorch 1.6**.
+    - OpenCV, NumPy, Pillow, tqdm, lpips, scikit-image and tensorboardX.
+
+- **Dataset**
+    - **GoPro dataset** can be downloaded from this [link](https://drive.google.com/file/d/1y4wvPdOG3mojpFCHTqLgriexhbjoWVkK/view).
+    - **Synthetic noisy images with Gaussian noise for testing** can be downloaded from this [link](https://pan.baidu.com/s/1eA8r5QoX0cLXSfikk6XlQw?pwd=vagc). Please decompress the files according to the commands provided in its `readme.txt`.
+    - **Synthetic noisy images with Poisson noise for testing** can be downloaded from this [link](https://pan.baidu.com/s/1tCCMxk7mlIk-27RD2_8GaA?pwd=fdw6). Please decompress the files according to the commands provided in its `readme.txt`.
+    - The directory structure of the dataset:
+
+        ```
+        GOPRO_Large 
+        │
+        └───train
+        │        GOPR0372_07_00
+        │        GOPR0372_07_01
+        |        ...  
+        │   
+        └───test
+        │        GOPR0372_07_00
+        │        GOPR0372_07_01
+        |        ... 
+        │   
+        └───test_noise_gauss5_50
+        │        GOPR0372_07_00
+        │        GOPR0372_07_01
+        |        ... 
+        │   
+        └───test_noise_poisson5_50
+        │        GOPR0372_07_00
+        │        GOPR0372_07_01
+        |        ... 
+        ```
+   
+
+    
+
+## 3. Quick Start
+
+### 3.1 Pre-trained models
+
+- For Gaussian noise, we provide the pre-trained models in the `./ckpt/selfir_gauss_noise/` folder.
+
+- For Poisson noise, we provide the pre-trained models in the `./ckpt/selfir_poisson_noise/` folder.
+
+
+### 3.2 Training
+
+- Modify `dataroot`, `type` and `name` in `train.sh` and then run:
+
+    [`sh train.sh`](train.sh)
+
+### 3.3 Testing
+
+- Modify `dataroot`, `type`, `name` and `iter` in `test.sh` and then run:
+
+    [`sh test.sh`](test.sh)
+
+### 3.4 Note
+
+- You can specify which GPU to use by `--gpu_ids`, e.g., `--gpu_ids 0,1`, `--gpu_ids 3`, `--gpu_ids -1` (for CPU mode). In the default setting, all GPUs are used.
+- You can refer to [options](./options/base_options.py) for more arguments.
+
+
+## 4. Citation
+If you find it useful in your research, please consider citing:
+
+    @inproceedings{SelfIR,
+        title={Self-Supervised Image Restoration with Blurry and Noisy Pairs},
+        author={Zhang, Zhilu and Xu, Rongjian and Liu, Ming and Yan, Zifei and Zuo, Wangmeng},
+        booktitle={NeurIPS},
+        year={2022}
+    }
+
+## 5. Acknowledgement
+
+This repo is built upon the framework of [CycleGAN](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix), and we borrow some code from [Neighbor2Neighbor](https://github.com/TaoHuang2018/Neighbor2Neighbor), thanks for their excellent work!

ckpt/selfir_gauss_noise/UNET_model_1.pth

@@ -0,0 +1,57 @@
+import importlib
+import torch.utils.data
+from data.base_dataset import BaseDataset
+
+
+def find_dataset_using_name(dataset_name, split='train'):
+    dataset_filename = "data." + dataset_name + "_dataset"
+    datasetlib = importlib.import_module(dataset_filename)
+
+    dataset = None
+    target_dataset_name = dataset_name.replace('_', '') + 'dataset'
+    for name, cls in datasetlib.__dict__.items():
+        if name.lower() == target_dataset_name.lower() \
+           and issubclass(cls, BaseDataset):
+            dataset = cls
+
+    if dataset is None:
+        raise NotImplementedError("In %s.py, there should be a subclass of "
+                        "BaseDataset with class name that matches %s in "
+                        "lowercase." % (dataset_filename, target_dataset_name))
+    return dataset
+
+
+def create_dataset(dataset_name, split, opt):
+    data_loader = CustomDatasetDataLoader(dataset_name, split, opt)
+    dataset = data_loader.load_data()
+    return dataset
+
+
+class CustomDatasetDataLoader():
+    def __init__(self, dataset_name, split, opt):
+        self.opt = opt
+        dataset_class = find_dataset_using_name(dataset_name, split)
+        self.dataset = dataset_class(opt, split, dataset_name)
+        self.imio = self.dataset.imio
+        print("dataset [%s(%s)] created" % (dataset_name, split))
+        self.dataloader = torch.utils.data.DataLoader(
+            self.dataset,
+            batch_size=opt.batch_size if split=='train' else 1,
+            shuffle=opt.shuffle and split=='train',
+            num_workers=int(opt.num_dataloader), 
+            drop_last=opt.drop_last)
+
+    def load_data(self):
+        return self
+
+    def __len__(self):
+        """Return the number of data in the dataset"""
+        return min(len(self.dataset), self.opt.max_dataset_size)
+
+    def __iter__(self):
+        """Return a batch of data"""
+        for i, data in enumerate(self.dataloader):
+            if i * self.opt.batch_size >= self.opt.max_dataset_size:
+                break
+            yield data
+

ckpt/selfir_poisson_noise/UNET_model_1.pth

@@ -0,0 +1,19 @@
+import torch.utils.data as data
+from abc import ABC, abstractmethod
+
+
+class BaseDataset(data.Dataset, ABC):
+    def __init__(self, opt, split, dataset_name):
+        self.opt = opt
+        self.split = split
+        self.root = opt.dataroot
+        self.dataset_name = dataset_name.lower()
+
+    @abstractmethod
+    def __len__(self):
+        return 0
+
+    @abstractmethod
+    def __getitem__(self, index):
+        pass
+ 

data/__init__.py

@@ -0,0 +1,149 @@
+import numpy as np
+import os
+from data.base_dataset import BaseDataset
+from .imlib import imlib
+from multiprocessing.dummy import Pool
+from tqdm import tqdm
+from util.util import augment
+import random
+
+
+# GoPro dataset
+class GoProDataset(BaseDataset):
+	def __init__(self, opt, split='train', dataset_name='GoPro'):
+		super(GoProDataset, self).__init__(opt, split, dataset_name)
+
+		if self.root == '':
+			rootlist = ['/Data/dataset/GOPRO_Large/']
+			for root in rootlist:
+				if os.path.isdir(root):
+					self.root = root
+					break
+
+		self.batch_size = opt.batch_size
+		self.patch_size = opt.patch_size
+		self.mode = opt.mode  # RGB, Y or L=
+		self.imio = imlib(self.mode, lib=opt.imlib)
+		self.names, self.blur_dirs, self.gt_dirs = self._get_image_dir(self.root, split)
+
+		if split == 'train':
+			self._getitem = self._getitem_train
+			self.len_data = 500 * 16 # 500 * self.batch_size
+		elif split == 'val':
+			self._getitem = self._getitem_test
+			self.len_data = len(self.names)
+		elif split == 'test': 
+			self._getitem = self._getitem_test
+			self.len_data = len(self.names)
+		else:
+			raise ValueError
+
+		self.blur_images = [0] * len(self.names)
+		self.gt_images = [0] * len(self.names)
+		read_images(self)
+
+	def __getitem__(self, index):
+		return self._getitem(index)
+
+	def __len__(self):
+		return self.len_data
+
+	def _getitem_train(self, idx):
+		idx = idx % len(self.names)
+
+		blur_img = self.blur_images[idx]
+		gt_img = self.gt_images[idx]
+		blur_img, gt_img = self._crop_patch(blur_img, gt_img)
+
+		blur_img, gt_img = augment(blur_img, gt_img)
+
+		blur_img = np.float32(blur_img) / 255 
+		gt_img = np.float32(gt_img) / 255
+
+		return {'gt_noise': gt_img,
+				'blur_img': blur_img,
+				'gt_img': gt_img,
+				'fname': self.names[idx]}
+
+	def _getitem_test(self, idx):
+
+		blur_img = self.blur_images[idx]
+		gt_img = self.gt_images[idx]
+
+		blur_img = np.float32(blur_img) / 255 
+		gt_img = np.float32(gt_img) / 255
+
+		noise_root = self.gt_dirs[idx].replace('sharp', 'npy')
+		noise_root = noise_root.replace('test', 'test_noise_' + self.opt.noisetype)
+		noise_file = noise_root[:-3] + 'npy'
+		gt_noise = np.float32(np.load(noise_file, allow_pickle=True))
+
+		return {'gt_noise': gt_noise,
+				'blur_img': blur_img,
+				'gt_img': gt_img,
+				'fname': self.names[idx]}
+
+	def _crop_patch(self, blur, gt):
+		ih, iw = blur.shape[-2:]
+		p = self.patch_size
+		pw = random.randrange(0, iw - p + 1)
+		ph = random.randrange(0, ih - p + 1)
+		return blur[..., ph:ph+p, pw:pw+p], \
+			   gt[..., ph:ph+p, pw:pw+p]
+
+	def _get_image_dir(self, dataroot, split=None):
+		blur_dirs = []
+		gt_dirs = [] 
+		image_names = []
+
+		if split == 'train' or split == 'test':
+			for scene_file in os.listdir(dataroot + split +  '/'): 
+				for image_file in os.listdir(dataroot + split +  '/' + scene_file + '/sharp/'):  
+					image_names.append(scene_file + '-' + image_file)
+					blur_dirs.append(dataroot + split +  '/' + scene_file + '/blur_gamma/' + image_file)
+					gt_dirs.append(dataroot + split +  '/' + scene_file + '/sharp/' + image_file)
+		elif split == 'val':
+			for scene_file in os.listdir(dataroot + 'test/'):
+				for image_file in os.listdir(dataroot + 'test/' + scene_file + '/sharp/'):  
+					image_names.append(scene_file + '-' + image_file)
+					blur_dirs.append(dataroot + 'test/' + scene_file + '/blur_gamma/' + image_file)
+					gt_dirs.append(dataroot + 'test/' + scene_file + '/sharp/' + image_file)
+					break
+		else:
+			raise ValueError
+
+		image_names = sorted(image_names)
+		blur_dirs = sorted(blur_dirs)
+		gt_dirs = sorted(gt_dirs)
+		
+		return image_names, blur_dirs, gt_dirs
+
+
+def iter_obj(num, objs):
+	for i in range(num):
+		yield (i, objs)
+
+def imreader(arg):
+	i, obj = arg
+	for _ in range(3):
+		try:
+			obj.blur_images[i] = obj.imio.read(obj.blur_dirs[i])
+			obj.gt_images[i] = obj.imio.read(obj.gt_dirs[i])
+			failed = False
+			break
+		except:
+			failed = True
+	if failed: print('%s fails!' % obj.names[i])
+
+def read_images(obj):
+	# may use `from multiprocessing import Pool` instead, but less efficient and
+	# NOTE: `multiprocessing.Pool` will duplicate given object for each process.
+	print('Starting to load images via multiple imreaders')
+	pool = Pool() # use all threads by default
+	for _ in tqdm(pool.imap(imreader, iter_obj(len(obj.names), obj)), total=len(obj.names)):
+		pass
+	pool.close()
+	pool.join()
+
+if __name__ == '__main__':
+	pass