[update] update readme
This commit is contained in:
LiuYuanchi
parent
353c4f7f86
commit
887cef5ebe
191
README.md
191
README.md
|
|
@ -1,96 +1,127 @@
|
|||
# Noise2Same
|
||||
Official TensorFlow implementation for the [paper](https://arxiv.org/abs/2010.11971) presented on NeurIPS 2020 titled
|
||||
"*Noise2Same: Optimizing A Self-Supervised Bound for Image Denoising*".
|
||||
# Noise2SamePro
|
||||
README
|
||||
|
||||
<img src="./figures/cover.jpg" width="80%">
|
||||
<img src="./figures/visual.jpg" width="80%">
|
||||
## 项目介绍
|
||||
本项目是基于自监督去噪模型[Noise2Same](https://github.com/divelab/Noise2Same)的改进，通过优化损失函数，提升了模型收敛速度和去噪效果.
|
||||
|
||||
## Environment Requirements
|
||||
- jupyter
|
||||
- python == 3.7.2
|
||||
- tensorflow >=1.10 & <=1.15
|
||||
## 使用效果
|
||||
去噪效果
|
||||
|
||||
|
||||
## 模型特点
|
||||
### 介绍
|
||||
本模型依赖于Noise2Same(下称N2S)模型的训练方法，在训练一定次数的N2S上继续训练得到Noise2SamePro(下称N2SP)，使得其收敛速度更快，图片去噪效果更好。
|
||||
与N2S模型的对比
|
||||
|
||||
|
||||
| | 方法 | PSNR |
|
||||
| :----: | :----: | :----: |
|
||||
| 传统方法 | Input | 20.19 |
|
||||
| | NLM | 22.73 |
|
||||
| 有监督 | Noise2True | 29.06 |
|
||||
| | Noise2Noise | 28.86 |
|
||||
| 自监督 | Noise2Same | 27.09 |
|
||||
| | Noise2SamePro | 27.10 |
|
||||
|
||||
### 数据集
|
||||
我们使用了N2S模型提供的[数据集](https://drive.google.com/drive/folders/1VYMo1OoaGxoOLNx6-qIt2Wg03lsZw_kA?usp=sharing),有关数据集构建和数据源的更多详细信息可以在 [Denoising_data](https://github.com/divelab/Noise2Same/blob/main/Denoising_data) 下找到.
|
||||
|
||||
### 检查点
|
||||
我们提供基于上述数据集训练9万轮次的检查点，请点击[这里](https://drive.google.com/drive/folders/1SL7Bx1TZgj8Ns4BJppA_PalCSY-J4L1R?usp=drive_link)下载.
|
||||
|
||||
## 使用方法
|
||||
### 环境
|
||||
```
|
||||
python==3.7.2
|
||||
```
|
||||
- tensorflow==1.15
|
||||
- scipy
|
||||
- skimage
|
||||
- scikit-image
|
||||
- tifffile
|
||||
- gdown
|
||||
- opencv-python
|
||||
- numpy
|
||||
- matplotlib
|
||||
- PIL
|
||||
|
||||
## Usage
|
||||
|
||||
### To reproduce our results
|
||||
|
||||
#### Dataset and model checkpoint download
|
||||
We uploaded the datasets used in our experiments and the model checkpoint files to the google drive [here](https://drive.google.com/drive/folders/1VYMo1OoaGxoOLNx6-qIt2Wg03lsZw_kA?usp=sharing). You can download the files and put them in the folders ``Denoising_data`` and ``trained_models``. More details about the dataset construction and the source of data can be found under [Denoising_data](./Denoising_data).
|
||||
|
||||
We have provided four examples in Jupyter Notebook that can reproduce our results in the paper. Once you have downloaded the dataset (and the pretrained chechpoints if you want to skip training), you can simply go through the notebooks for reproduction.
|
||||
|
||||
### To train, evaluate and predict with your own datasets
|
||||
You can follow the examples in Jupyter Notebook for denoising with RGB images, grayscale images and 3D images.
|
||||
|
||||
#### To be specific, the following code is used to build the model.
|
||||
### 目录树
|
||||
```
|
||||
from models import Noise2Same
|
||||
model = Noise2Same(model_dir, model_name, dimension, in_channels)
|
||||
```
|
||||
where ``model_dir`` and ``model_name`` will specify the path to your checkpoint files, ``dimension`` refers to the dimension of image *(2 or 3)* and ``in_channels`` refers to the number of channels of input images.
|
||||
|
||||
#### The following code is used for **training**.
|
||||
```
|
||||
model.train(X, patch_size, validation=X_val, batch_size, steps)
|
||||
```
|
||||
where ``X`` and ``X_val`` are the noisy images for training/validation of shape ``[n_samples, width, length, n_channels]`` and of type ``float32``, ``patch_size`` specify the size to crop input images to training patches. Note that the input image should be **normalized** before input for training.
|
||||
|
||||
#### The following codes are for **prediction**.
|
||||
|
||||
- For prediction of single image,
|
||||
```
|
||||
model.predict(img[, im_mean, im_std])
|
||||
```
|
||||
where ``img`` is the noisy image for prediction, ``im_mean`` and ``im_std`` are the mean and standard deviation. If ``im_mean`` and ``im_std`` are not specified, it will use ``img.mean()`` and ``img.std()`` by default.
|
||||
|
||||
- For prediction of batched images (and you have enough GPU memory),
|
||||
```
|
||||
model.batch_predict(images.astype('float32'), batch_size[, im_mean, im_std])
|
||||
```
|
||||
|
||||
- For extremely large images, e.g. CARE 3D images,
|
||||
```
|
||||
model.crop_predict(image, crop_size, overlap[, im_mean, im_std])
|
||||
```
|
||||
|
||||
### Use Noise2Same under other frameworks
|
||||
You can follow the pseudocode below to build the Noise2Same model.
|
||||
|
||||
Given the noisy images ``images``, the masked noisy images ``masked_images`` and masking map ``mask`` with masked locations being 1 and other 0,
|
||||
Noise2SamePro
|
||||
├─ test.py *
|
||||
├─ test_pro.py *
|
||||
├─ network_configure.py
|
||||
├─ models.py *
|
||||
├─ requirements.txt
|
||||
├─ basic_ops.py
|
||||
├─ train.py *
|
||||
├─ train_pro.py *
|
||||
├─ utils
|
||||
│ ├─ train_utils.py
|
||||
│ ├─ evaluation_utils.py
|
||||
│ └─ predict_utils.py
|
||||
├─ README.md
|
||||
├─ network.py
|
||||
├─ resnet_module.py
|
||||
├─ Denoising_data ------------------- 默认的数据集文件夹
|
||||
├─ trained_model -------------------- 默认的模型存放位置
|
||||
└─ test_single.py
|
||||
|
||||
```
|
||||
net = YourNetwork()
|
||||
# The two net() below should share their weights
|
||||
out_raw = net(images)
|
||||
out_masked = net(masked_images)
|
||||
您需要关注的是标记有`*`的文件/文件夹，后文会介绍通过修改这些文件进行自定义训练的方法.本项目源代码中同时含有N2S模型与N2SP模型，其中`test.py`,`train.py`是有关N2S模型的训练与测试部分；`test_pro.py`,`train_pro.py`是有关N2SP的训练与测试部分.
|
||||
|
||||
l_rec = reduce_mean((out_raw - images)^2)
|
||||
l_inv = reduce_sum((out_raw - out_masked)^2 * mask) / reduce_sum(mask)
|
||||
loss = l_rec + 2 * sqrt(l_inv)
|
||||
下文介绍N2SP的测试与训练方法，N2S与之类似(您也可参照N2S原仓库).
|
||||
|
||||
### 使用
|
||||
在`test_single.py`中，请修改如下部分进行单张图片的去噪
|
||||
```python
|
||||
picture = 'man/' # Adjust path of the picture you want to test
|
||||
model_dir = 'N2S_PRO' # Adjust your model path
|
||||
test_single('test_single/' + picture + 'original_image.png', model_dir, 'test_single/' + picture + 'denoised_image.png')
|
||||
```
|
||||
执行以下命令并等待即可
|
||||
```shell
|
||||
python test_single.py
|
||||
```
|
||||
### 测试
|
||||
如果您使用我们提供的检查点和默认路径，请将其放在`trained_model/`文件夹下(没有则请新建)
|
||||
确保`test_pro.py`文件中检查点路径与数据集路径正确.
|
||||
```python
|
||||
model_dir = 'N2S_PRO-90000' # Adjust your model path
|
||||
data_dir = 'Denoising_data/test/'
|
||||
model = Noise2Same('trained_models/', model_dir, dim=2, in_channels=1)
|
||||
```
|
||||
执行以下命令并等待即可
|
||||
```shell
|
||||
python test_pro.py
|
||||
```
|
||||
|
||||
## Reference
|
||||
```
|
||||
@inproceedings{xie2020noise2same,
|
||||
author = {Xie, Yaochen and Wang, Zhengyang and Ji, Shuiwang},
|
||||
title = {Noise2{S}ame: Optimizing A Self-Supervised Bound for Image Denoising},
|
||||
booktitle = {Advances in Neural Information Processing Systems},
|
||||
pages = {20320--20330},
|
||||
volume = {33},
|
||||
year = {2020}
|
||||
}
|
||||
### 训练
|
||||
N2SP的训练方法是：
|
||||
|
||||
> 对于T次训练，采取`(T-a)`次对于N2S的训练和`a`次对N2SP的训练
|
||||
|
||||
确保路径正确后，修改`train_pro.py`
|
||||
```python
|
||||
model_dir = 'N2S_PRO-8000' # Set model checkpoints save path
|
||||
steps = 8000 # Set training steps
|
||||
|
||||
sgm_loss = 1 # the default sigma is 1
|
||||
model = Noise2Same('trained_models/', model_dir, dim=2, in_channels=1, lmbd=2*sgm_loss)
|
||||
model.train(X[..., None], patch_size=[64, 64], validation=X_val[..., None], batch_size=64, steps=steps-500)
|
||||
model = Noise2SamePro('trained_models/', model_dir, dim=2, in_channels=1)
|
||||
model.train(X[..., None], patch_size=[64, 64], validation=X_val[..., None], batch_size=64, steps=steps)
|
||||
|
||||
```
|
||||
|
||||
## Web Demo
|
||||
You can create a web-based demo to run inference by running the `demo.py` file, which uses the `gradio` Python library.
|
||||
在上面的数值中，`steps = 8000`是总步数，`sgm_loss = 1`是N2S模型使用的参数(默认为1),`steps=steps-500`中的500是N2SP的训练次数
|
||||
即8000为 上述`T`;500为上述`a`.训练过程是进行7500次N2S的预训练，然后进行500次的附加训练.
|
||||
|
||||
Here is a live demo: https://gradio.app/g/Noise2Same
|
||||
调整完毕后，执行以下命令并等待即可
|
||||
```shell
|
||||
python train_pro.py
|
||||
```
|
||||
|
||||
The live demo uses the model pre-trained on 20,000 noisy images generated from ImageNet ILSVRC2012 validation dataset.
|
||||
### 其他事项
|
||||
|
||||
|
||||
|
||||
We thank [Abubakar Abid](https://github.com/abidlabs) for building this awesome web demo for us!
|
||||
## 贡献者
|
||||
本项目是天津大学2021图像处理结课作业的去噪部分，由郭骐瑞，刘原驰，罗嘉杰，周欣宇，陈嘉强，李进共同完成
|
||||
Loading…
Reference in New Issue