forked from huawei/mindspore2022
add dncc to model zoo
This commit is contained in:
parent
303d4857e8
commit
288e51b623
|
|
@ -0,0 +1,139 @@
|
|||
|
||||
### 目录
|
||||
|
||||
[TOC]
|
||||
|
||||
### DnCNN描述
|
||||
|
||||
DnCNN是一个使用FCN处理图像降噪任务的模型， 本项目是图像去躁模型DnCNN在mindspore上的复现。
|
||||
论文\: Zhang, K., Zuo, W., Chen, Y., Meng, D., & Zhang, L. (2017). Beyond a gaussian denoiser: Residual learning of deep cnn for image denoising. IEEE transactions on image processing, 26(7), 3142-3155.
|
||||
|
||||
### 模型结构
|
||||
|
||||
网络由N层convolution block组成。其中第一层是conv加reLU，中间n-2层是conv+BN+ReLU，最后一层是单独的conv
|
||||
|
||||
### 数据集
|
||||
|
||||
训练集DnCNN-S、DnCNN-B使用BSD500中的400张图片
|
||||
DnCNN-3使用BSD500中的200张图片和T91中的91张图片
|
||||
|
||||
测试集包括BDS68，Set5， Set14， clasic5， live1等
|
||||
|
||||
### 环境要求
|
||||
|
||||
mindspore=1.1
|
||||
skimage=0.18.1
|
||||
numpy
|
||||
PIL
|
||||
opencv
|
||||
argparse
|
||||
|
||||
### 快速入门
|
||||
|
||||
通过官方网站安装MindSpore后，您可以按照如下步骤进行训练和评估：
|
||||
|
||||
```bash
|
||||
# 训练示例
|
||||
python train.py --dataset_path=/path/to/training/data --model_type DnCNN-S --ckpt-prefix=DnCNN-S_25noise --noise_level=25
|
||||
# 或者
|
||||
sh ./scripts/run_train_gpu.sh /path/to/training/data DnCNN-S DnCNN-S_25noise 25
|
||||
|
||||
# 评估示例
|
||||
python eval.py --dataset_path=/path/to/test/data --ckpt_path=./ckpt/DnCNN-S-50_1800.ckpt --model_type=DnCNN-S --noise_level=25 --noise_type=denoise
|
||||
# 或者
|
||||
sh ./scripts/run_eval_gpu.sh /path/to/test/data ./ckpt/DnCNN-S-50_1800.ckpt DnCNN-S 25 denoise
|
||||
```
|
||||
|
||||
### 脚本说明
|
||||
|
||||
├── readme.md
|
||||
├── scripts
|
||||
│ ├── run_eval_gpu.sh //训练shell脚本
|
||||
│ └── run_train_gpu.sh //评估shell脚本
|
||||
├── src
|
||||
│ ├── dataset.py //数据读取
|
||||
│ └──model.py //模型定义
|
||||
├── eval.py //评估脚本
|
||||
├── export.py //导出模型
|
||||
└── train.py //训练脚本
|
||||
|
||||
### 训练过程
|
||||
|
||||
可通过`train.py`脚本中的参数修改训练行为。`train.py`脚本中的参数如下：
|
||||
|
||||
#### 训练参数
|
||||
|
||||
--dataset_path 训练数据路径
|
||||
--model_type 模型类型 = ['DnCNN-S', 'DnCNN-B', 'DnCNN-3']
|
||||
--ckpt-prefix 检查点前缀
|
||||
--noise_level 噪音等级
|
||||
--batch_size 批次大小
|
||||
--lr 学习率
|
||||
--epoch_num 轮次数
|
||||
|
||||
#### 默认训练参数
|
||||
|
||||
optimizer=adam
|
||||
learning rate=0.001
|
||||
batch_size=128
|
||||
weight_decay=0.0001
|
||||
epoch=50
|
||||
|
||||
#### 训练
|
||||
|
||||
只有DnCNN-S 需要指定noise_level
|
||||
|
||||
```python
|
||||
python train.py --dataset_path=/path/to/training/data --model_type=DnCNN-S --ckpt-prefix=DnCNN-S_25noise --noise_level=25
|
||||
python train.py --dataset_path=/path/to/training/data --model_type=DnCNN-B --ckpt-prefix=DnCNN-B
|
||||
python train.py --dataset_path=/path/to/training/data --model_type=DnCNN-3 --ckpt-prefix=DnCNN-3
|
||||
```
|
||||
|
||||
在ckpt文件夹下保存检查点
|
||||
|
||||
### 评估过程
|
||||
|
||||
评估需要通过命令行提供以下参数：
|
||||
--dataset_path 数据路径
|
||||
--ckpt_path 检查点路径
|
||||
--model_type 模型类型
|
||||
--noise_type 噪音类型， 通过noise_type选择图像测试噪音的类型：["denoise", "super-resolution","jpeg-deblock"]
|
||||
--noise_level 噪音等级：对应三种noise type的强度，噪音sigma/下采样上采样scale/jpeg压缩quality
|
||||
|
||||
ex:
|
||||
|
||||
```python
|
||||
python eval.py --dataset_path=/path/to/test/data --ckpt_path=./ckpt/DnCNN-B-50_3000.ckpt --model_type=DnCNN-B --noise_level=50 --noise_type=denoise
|
||||
```
|
||||
|
||||
### 模型描述
|
||||
|
||||
#### 训练准确率结果
|
||||
|
||||
| 参数 | GPU |
|
||||
| ------------- | -------------------------------------------------------- |
|
||||
| 模型版本 | DnCNN-S |
|
||||
| 资源 | Nvidia V100 |
|
||||
| mindspore版本 | mindspore 1.1 |
|
||||
| 数据集 | Berkeley Segmentation Datase |
|
||||
| 轮次 | 50 |
|
||||
| 输出 | noise残差 |
|
||||
| 性能 | 在BSD68测试，PSNR=32.92(σ=15)， 31.73(σ=25)，30.59(σ=50) |
|
||||
|
||||
#### 训练性能结果
|
||||
|
||||
| 参数 | GPU |
|
||||
| ------------- | ---------------------------------- |
|
||||
| 模型版本 | DnCNN-S |
|
||||
| 资源 | Nvidia V100 |
|
||||
| mindspore版本 | mindspore 1.1 |
|
||||
| 训练参数 | lr 0.001, batch_size 128, epoch 50 |
|
||||
| 优化器 | adam |
|
||||
| 损失函数 | MSE |
|
||||
| 输出 | noise残差 |
|
||||
| 速度 | 320ms/batch |
|
||||
| 总时长 | 7h9min |
|
||||
| 检查点 | 6.38M |
|
||||
|
||||
|
||||
|
||||
|
|
@ -0,0 +1,146 @@
|
|||
#!/usr/bin/env python3
|
||||
# Copyright 2021 Huawei Technologies Co., Ltd
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ============================================================================
|
||||
|
||||
import os
|
||||
import io
|
||||
import argparse
|
||||
import glob
|
||||
|
||||
import PIL
|
||||
import numpy as np
|
||||
import cv2
|
||||
import skimage.metrics
|
||||
import mindspore
|
||||
import mindspore.dataset as ds
|
||||
from mindspore import context
|
||||
from mindspore.train.serialization import load_checkpoint, load_param_into_net
|
||||
import mindspore.dataset.transforms.c_transforms as C
|
||||
from src.model import DnCNN
|
||||
|
||||
class DnCNN_eval_Dataset():
|
||||
def __init__(self, dataset_path, task_type, noise_level):
|
||||
self.im_list = []
|
||||
self.im_list.extend(glob.glob(os.path.join(dataset_path, "*.png")))
|
||||
self.im_list.extend(glob.glob(os.path.join(dataset_path, "*.bmp")))
|
||||
self.im_list.extend(glob.glob(os.path.join(dataset_path, "*.jpg")))
|
||||
self.task_type = task_type
|
||||
self.noise_level = noise_level
|
||||
|
||||
def __getitem__(self, i):
|
||||
img = cv2.imread(self.im_list[i], 0)
|
||||
|
||||
if self.task_type == "denoise":
|
||||
noisy = self.add_noise(img, self.noise_level)
|
||||
elif self.task_type == "super-resolution":
|
||||
h, w = img.shape
|
||||
noisy = cv2.resize(img, (int(w/self.noise_level), int(h/self.noise_level)))
|
||||
noisy = cv2.resize(noisy, (w, h))
|
||||
elif self.task_type == "jpeg-deblock":
|
||||
noisy = self.jpeg_compression(img, self.noise_level)
|
||||
|
||||
#add channel dimension
|
||||
noisy = noisy[np.newaxis, :, :]
|
||||
noisy = noisy / 255.0
|
||||
return noisy, img
|
||||
|
||||
def __len__(self):
|
||||
return len(self.im_list)
|
||||
|
||||
def add_noise(self, im, sigma):
|
||||
gauss = np.random.normal(0, sigma, im.shape)
|
||||
noisy = im + gauss
|
||||
noisy = np.clip(noisy, 0, 255)
|
||||
noisy = noisy.astype('float32')
|
||||
return noisy
|
||||
|
||||
def jpeg_compression(self, img, severity):
|
||||
im_pil = PIL.Image.fromarray(img)
|
||||
output = io.BytesIO()
|
||||
im_pil.save(output, 'JPEG', quality=severity)
|
||||
im_pil = PIL.Image.open(output)
|
||||
img_np = np.asarray(im_pil)
|
||||
return img_np
|
||||
|
||||
|
||||
def create_eval_dataset(data_path, task_type, noise_level, batch_size=1):
|
||||
# define dataset
|
||||
dataset = DnCNN_eval_Dataset(data_path, task_type, noise_level)
|
||||
dataloader = ds.GeneratorDataset(dataset, ["noisy", "clear"])
|
||||
# apply map operations on images
|
||||
dataloader = dataloader.map(input_columns="noisy", operations=C.TypeCast(mindspore.float32))
|
||||
dataloader = dataloader.map(input_columns="clear", operations=C.TypeCast(mindspore.uint8))
|
||||
dataloader = dataloader.batch(batch_size, drop_remainder=False)
|
||||
return dataloader
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="DnCNN")
|
||||
parser.add_argument("--dataset_path", type=str, default="/code/12imgs-TestingSet/", help='eval image path')
|
||||
parser.add_argument('--ckpt_path', type=str, default=False, help='trained ckpt')
|
||||
parser.add_argument('--model_type', type=str, default='DnCNN-S', \
|
||||
choices=['DnCNN-S', 'DnCNN-B', 'DnCNN-3'], help='type of DnCNN')
|
||||
parser.add_argument('--noise_type', type=str, default=False, \
|
||||
choices=["denoise", "super-resolution", "jpeg-deblock"], help='trained ckpt')
|
||||
parser.add_argument('--noise_level', type=int, default=False, help='trained ckpt')
|
||||
args = parser.parse_args()
|
||||
|
||||
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
|
||||
ds_eval = create_eval_dataset(args.dataset_path, args.noise_type, args.noise_level, batch_size=1)
|
||||
print("evaluation image number:", ds_eval.get_dataset_size())
|
||||
|
||||
if args.model_type == 'DnCNN-S':
|
||||
network = DnCNN(1, num_of_layers=17)
|
||||
elif args.model_type == 'DnCNN-3' or args.model_type == 'DnCNN-B':
|
||||
network = DnCNN(1, num_of_layers=20)
|
||||
else:
|
||||
print("wrong model type")
|
||||
exit()
|
||||
|
||||
# load parameter to the network
|
||||
param_dict = load_checkpoint(args.ckpt_path)
|
||||
load_param_into_net(network, param_dict)
|
||||
|
||||
mean_psnr = 0
|
||||
mean_ssim = 0
|
||||
count = 0
|
||||
for data in ds_eval.create_dict_iterator():
|
||||
clear = data["clear"].asnumpy()
|
||||
#get denoised image
|
||||
residual = network(data["noisy"]).asnumpy() * 255
|
||||
noisy_img = data["noisy"].asnumpy() * 255
|
||||
denoised = np.clip(noisy - residual, 0, 255).astype("uint8")
|
||||
denoised = np.squeeze(denoised)
|
||||
clear = np.squeeze(clear)
|
||||
noisy_img = np.squeeze(noisy_img)
|
||||
if count == 0: #save example result
|
||||
cv2.imwrite("noisy.jpg", noisy_img.astype("uint8"))
|
||||
cv2.imwrite("denoised.jpg", denoised)
|
||||
cv2.imwrite("original.jpg", clear)
|
||||
|
||||
#calculate psnr
|
||||
mse = np.mean((clear - denoised)**2)
|
||||
psnr = 10*np.log10(255*255/mse)
|
||||
#calculate ssim
|
||||
ssim = skimage.metrics.structural_similarity(clear, denoised, data_range=255) #skimage 0.18
|
||||
|
||||
mean_psnr += psnr
|
||||
mean_ssim += ssim
|
||||
count += 1
|
||||
|
||||
mean_psnr = mean_psnr / count
|
||||
mean_ssim = mean_ssim / count
|
||||
print("mean psnr", mean_psnr)
|
||||
print("mean_ssim", mean_ssim)
|
||||
|
|
@ -0,0 +1,52 @@
|
|||
#!/usr/bin/env python3
|
||||
# Copyright 2021 Huawei Technologies Co., Ltd
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ============================================================================
|
||||
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
import mindspore as ms
|
||||
from mindspore import Tensor, load_checkpoint, load_param_into_net, export, context
|
||||
|
||||
from src.model import DnCNN
|
||||
|
||||
parser = argparse.ArgumentParser(description='DnCNN')
|
||||
parser.add_argument("--batch_size", type=int, default=1, help="batch size")
|
||||
parser.add_argument("--image_height", type=int, default=256, help="image_height")
|
||||
parser.add_argument("--image_width", type=int, default=256, help="image_width")
|
||||
parser.add_argument("--ckpt_file", type=str, required=True, help="Checkpoint file path.")
|
||||
parser.add_argument("--file_name", type=str, default="DnCNN", help="output file name.")
|
||||
parser.add_argument('--file_format', type=str, choices=["AIR", "ONNX", "MINDIR"], default='MINDIR', help='file format')
|
||||
parser.add_argument('--model_type', type=str, default='DnCNN-S', \
|
||||
choices=['DnCNN-S', 'DnCNN-B', 'DnCNN-3'], help='type of DnCNN')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
|
||||
if args.model_type == 'DnCNN-S':
|
||||
network = DnCNN(1, num_of_layers=17)
|
||||
elif args.model_type == 'DnCNN-3' or args.model_type == 'DnCNN-B':
|
||||
network = DnCNN(1, num_of_layers=20)
|
||||
else:
|
||||
print("wrong model type")
|
||||
exit()
|
||||
|
||||
param_dict = load_checkpoint(args.ckpt_file)
|
||||
load_param_into_net(network, param_dict)
|
||||
|
||||
input_arr = Tensor(np.ones([args.batch_size, 1, args.image_height, args.image_width]), ms.float32)
|
||||
export(network, input_arr, file_name=args.file_name, file_format=args.file_format)
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
mindspore-gpu==1.1.1
|
||||
numpy==1.17.0
|
||||
Pillow=6.2.2
|
||||
opencv-python==4.2.0.34
|
||||
scikit-image==0.18.1
|
||||
|
|
@ -0,0 +1,16 @@
|
|||
#!/bin/bash
|
||||
# Copyright 2021 Huawei Technologies Co., Ltd
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ============================================================================
|
||||
python eval.py --dataset_path=$1 --ckpt_path=$2 --model_type=$3 --noise_level=$4 --noise_type=$5
|
||||
|
|
@ -0,0 +1,21 @@
|
|||
#!/bin/bash
|
||||
# Copyright 2021 Huawei Technologies Co., Ltd
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ============================================================================
|
||||
if [ "$2" == "DnCNN-S" ]
|
||||
then
|
||||
python train.py --dataset_path=$1 --model_type=$2 --ckpt_prefix=$3 --noise_level=$4
|
||||
else
|
||||
python train.py --dataset_path=$1 --model_type=$2 --ckpt_prefix=$3
|
||||
fi
|
||||
|
|
@ -0,0 +1,182 @@
|
|||
#!/usr/bin/env python3
|
||||
# Copyright 2021 Huawei Technologies Co., Ltd
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ============================================================================
|
||||
|
||||
import os
|
||||
import random
|
||||
import glob
|
||||
import io
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
import PIL
|
||||
import mindspore
|
||||
import mindspore.dataset as ds
|
||||
import mindspore.dataset.transforms.c_transforms as C
|
||||
|
||||
def create_train_dataset(data_path, model_type, noise_level=25, batch_size=128):
|
||||
# define dataset
|
||||
if model_type == "DnCNN-S":
|
||||
dataset = DnCNN_train_Dataset(data_path, model_type, batch_size, noise_level, patch_shape=(40, 40))
|
||||
if model_type in ["DnCNN-B", "DnCNN-3"]:
|
||||
dataset = DnCNN_train_Dataset(data_path, model_type, batch_size, patch_shape=(50, 50))
|
||||
|
||||
print("total training patch numbers per epoch", len(dataset))
|
||||
dataloader = ds.GeneratorDataset(dataset, ["noisy", "gt"])
|
||||
# apply map operations on images
|
||||
dataloader = dataloader.map(input_columns="noisy", operations=C.TypeCast(mindspore.float32))
|
||||
dataloader = dataloader.map(input_columns="gt", operations=C.TypeCast(mindspore.float32))
|
||||
|
||||
# apply DatasetOps
|
||||
dataloader = dataloader.shuffle(buffer_size=10000)
|
||||
dataloader = dataloader.batch(batch_size, drop_remainder=True)
|
||||
dataloader = dataloader.repeat(1) #here 400 images as an epoch , on the paper 128x1600 patches as a epoch
|
||||
return dataloader
|
||||
|
||||
class DnCNN_train_Dataset():
|
||||
def __init__(self, dataset_path, model_type, batch_size=128, noise_level=25, \
|
||||
image_shape=(180, 180), patch_shape=(50, 50)):
|
||||
#DnCNN-S/B uses 200 training and 200 test images in BDS500 as training set
|
||||
#DnCNN-3 uses 200 training images in BDS500 and T91 images as training set
|
||||
self.im_list = []
|
||||
self.im_list.extend(glob.glob(os.path.join(dataset_path, "*/*.jpg")))
|
||||
self.im_list.extend(glob.glob(os.path.join(dataset_path, "*/*.bmp")))
|
||||
self.im_list.extend(glob.glob(os.path.join(dataset_path, "*/*.png")))
|
||||
self.patch_shape = patch_shape
|
||||
self.image_shape = image_shape
|
||||
self.batch_size = batch_size
|
||||
self.model_type = model_type
|
||||
self.noise_level = noise_level
|
||||
self.scales = [0.8, 1, 1.25]
|
||||
self.sr_scales = [2, 3, 4]
|
||||
self.compression_range = (5, 100)
|
||||
self.sigma_range = (0, 55)
|
||||
|
||||
def __getitem__(self, i):
|
||||
i = i % len(self.im_list)
|
||||
img = cv2.imread(self.im_list[i], 0)
|
||||
#random scale
|
||||
scale = random.choice(self.scales)
|
||||
img_h = int(self.image_shape[1]*scale)
|
||||
img_w = int(self.image_shape[0]*scale)
|
||||
img = cv2.resize(img, (img_w, img_h))
|
||||
|
||||
#crop random patch
|
||||
start_w = random.randint(0, img_w - self.patch_shape[0])
|
||||
start_h = random.randint(0, img_h - self.patch_shape[1])
|
||||
patch = img[start_h:start_h+self.patch_shape[1], start_w:start_w+self.patch_shape[0]]
|
||||
|
||||
#random flip & rotation
|
||||
patch = self.data_augment(patch)
|
||||
|
||||
#add noise
|
||||
if self.model_type == "DnCNN-S":
|
||||
#add specific level of noise
|
||||
noisy = self.add_noise(patch, self.noise_level)
|
||||
elif self.model_type == "DnCNN-B":
|
||||
#add random level of noise
|
||||
sigma = random.uniform(*self.sigma_range)
|
||||
noisy = self.add_noise(patch, sigma)
|
||||
elif self.model_type == "DnCNN-3":
|
||||
#randomly add noise, down-up sampling blur and jpeg bloc
|
||||
noisy = self.add_random_type_noise(patch)
|
||||
else:
|
||||
print("wrong type")
|
||||
exit()
|
||||
noise = np.float32(noisy) - np.float32(patch)
|
||||
|
||||
#add channel dimension
|
||||
noisy = noisy[np.newaxis, :, :]
|
||||
noise = noise[np.newaxis, :, :]
|
||||
|
||||
#normalize
|
||||
noisy = noisy / 255.0
|
||||
noise = noise / 255.0
|
||||
|
||||
return noisy, noise
|
||||
|
||||
def __len__(self):
|
||||
#To build the same epoch size as original paper
|
||||
#on the paper, DnCNN-S has 1600 iteratisons per epoch
|
||||
#DnCNN-B has 3000 and DnCNN-3 has 8000
|
||||
if self.model_type == "DnCNN-S":
|
||||
epoch_size = self.batch_size * 1600
|
||||
elif self.model_type == "DnCNN-B":
|
||||
epoch_size = self.batch_size * 3000
|
||||
elif self.model_type == "DnCNN-3":
|
||||
epoch_size = self.batch_size * 8000
|
||||
else:
|
||||
print("wrong model type")
|
||||
exit()
|
||||
return epoch_size
|
||||
|
||||
def data_augment(self, patch):
|
||||
#random flip & rotation
|
||||
if random.random() < 0.5:
|
||||
patch = np.fliplr(patch)
|
||||
if random.random() < 0.5:
|
||||
patch = np.flipud(patch)
|
||||
alea = random.random()
|
||||
if alea > 0.25:
|
||||
patch = np.rot90(patch)
|
||||
elif alea > 0.5:
|
||||
patch = np.rot90(patch, k=2)
|
||||
elif alea > 0.75:
|
||||
patch = np.rot90(patch, k=3)
|
||||
return patch
|
||||
|
||||
def add_random_type_noise(self, patch):
|
||||
#random noise/down-up resampling/JPEG compression
|
||||
alea = random.random()
|
||||
if alea < 0.33: #generate noisy image
|
||||
sigma = random.uniform(*self.sigma_range)
|
||||
noisy = self.add_noise(patch, sigma)
|
||||
elif alea < 0.66: #generate blur image
|
||||
sr_scale = random.choice(self.sr_scales)
|
||||
noisy = cv2.resize(patch, (int(self.patch_shape[0]/sr_scale), int(self.patch_shape[1]/sr_scale)))
|
||||
noisy = cv2.resize(noisy, (int(self.patch_shape[0]), int(self.patch_shape[1])))
|
||||
else: #generate JPEG blocking image
|
||||
compression_level = random.randint(*self.compression_range)
|
||||
noisy = self.jpeg_compression(patch, compression_level)
|
||||
return noisy
|
||||
|
||||
def add_noise(self, im, sigma):
|
||||
gauss = np.random.normal(0, sigma, self.patch_shape)
|
||||
noisy = im + gauss
|
||||
noisy = np.clip(noisy, 0, 255)
|
||||
noisy = noisy.astype('uint8')
|
||||
return noisy
|
||||
|
||||
def jpeg_compression(self, img, quality):
|
||||
im_pil = PIL.Image.fromarray(img)
|
||||
output = io.BytesIO()
|
||||
im_pil.save(output, 'JPEG', quality=quality)
|
||||
im_pil = PIL.Image.open(output)
|
||||
img_np = np.asarray(im_pil)
|
||||
return img_np
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
#only for test
|
||||
test_dataset_path = "/code/BSR_bsds500/BSR/BSDS500/data/images/"
|
||||
ds_train = create_train_dataset(test_dataset_path, "DnCNN-S", batch_size=128)
|
||||
print("batch number:", ds_train.get_dataset_size())
|
||||
|
||||
for data in ds_train.create_dict_iterator():
|
||||
print(type(data))
|
||||
print(data["noisy"].shape)
|
||||
print(data["gt"].shape)
|
||||
print(type(data["noisy"]))
|
||||
break
|
||||
|
|
@ -0,0 +1,59 @@
|
|||
#!/usr/bin/env python3
|
||||
# Copyright 2021 Huawei Technologies Co., Ltd
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ============================================================================
|
||||
|
||||
import numpy as np
|
||||
import mindspore
|
||||
import mindspore.nn as nn
|
||||
from mindspore import context
|
||||
from mindspore import Tensor
|
||||
from mindspore.common.initializer import HeUniform
|
||||
|
||||
|
||||
class DnCNN(nn.Cell):
|
||||
def __init__(self, channels, num_of_layers=20):
|
||||
super(DnCNN, self).__init__()
|
||||
kernel_size = 3
|
||||
padding = 1
|
||||
features = 64
|
||||
layers = []
|
||||
layers.append(nn.Conv2d(channels, out_channels=features, kernel_size=kernel_size, \
|
||||
pad_mode='pad', padding=padding, has_bias=False, weight_init=HeUniform()))
|
||||
layers.append(nn.ReLU())
|
||||
|
||||
for _ in range(num_of_layers-2):
|
||||
layers.append(nn.Conv2d(features, out_channels=features, kernel_size=kernel_size, \
|
||||
pad_mode='pad', padding=padding, has_bias=False, weight_init=HeUniform()))
|
||||
layers.append(nn.BatchNorm2d(features))
|
||||
layers.append(nn.ReLU())
|
||||
|
||||
layers.append(nn.Conv2d(features, out_channels=channels, kernel_size=kernel_size, \
|
||||
pad_mode='pad', padding=padding, has_bias=False, weight_init=HeUniform()))
|
||||
self.dncnn = nn.SequentialCell(layers)
|
||||
|
||||
def construct(self, x):
|
||||
out = self.dncnn(x)
|
||||
return out
|
||||
|
||||
if __name__ == "__main__":
|
||||
#for test
|
||||
context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
|
||||
net = DnCNN(1, num_of_layers=17)
|
||||
a = Tensor(np.ones((2, 1, 40, 40)), mindspore.float32)
|
||||
output = net(a)
|
||||
print(output)
|
||||
print(type(output))
|
||||
np_out = output.asnumpy()
|
||||
print(type(np_out))
|
||||
|
|
@ -0,0 +1,92 @@
|
|||
#!/usr/bin/env python3
|
||||
# Copyright 2021 Huawei Technologies Co., Ltd
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ============================================================================
|
||||
|
||||
import argparse
|
||||
import datetime
|
||||
import mindspore.nn as nn
|
||||
from mindspore import context
|
||||
from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor, LearningRateScheduler
|
||||
from mindspore.train import Model
|
||||
from mindspore.train.callback import Callback
|
||||
|
||||
from src.dataset import create_train_dataset
|
||||
from src.model import DnCNN
|
||||
|
||||
|
||||
class BatchAverageMSELoss(nn.Cell):
|
||||
def __init__(self, batch_size):
|
||||
super(BatchAverageMSELoss, self).__init__()
|
||||
self.batch_size = batch_size
|
||||
self.sumMSELoss = nn.MSELoss(reduction='sum')
|
||||
|
||||
def construct(self, logits, labels):
|
||||
#equation 1 on the paper
|
||||
loss = self.sumMSELoss(logits, labels) / self.batch_size / 2
|
||||
return loss
|
||||
|
||||
class Print_info(Callback):
|
||||
def epoch_end(self, run_context):
|
||||
cb_params = run_context.original_args()
|
||||
print(datetime.datetime.now(), "end epoch", cb_params.cur_epoch_num)
|
||||
|
||||
def learning_rate_function(lr, cur_step_num):
|
||||
if cur_step_num % 40000 == 0:
|
||||
lr = lr*0.8
|
||||
print("current lr: ", str(lr))
|
||||
return lr
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="DnCNN")
|
||||
parser.add_argument("--dataset_path", type=str, default="/code/BSR_bsds500/BSR/BSDS500/data/images/", \
|
||||
help='training image path')
|
||||
parser.add_argument("--batch_size", type=int, default=128, help='training batch size')
|
||||
parser.add_argument('--lr', type=float, default=0.001, help='learning rate')
|
||||
parser.add_argument('--weight_decay', type=float, default=0.0001, help='weight_decay')
|
||||
parser.add_argument('--model_type', type=str, default='DnCNN-S', \
|
||||
choices=['DnCNN-S', 'DnCNN-B', 'DnCNN-3'], help='type of DnCNN')
|
||||
parser.add_argument('--noise_level', type=int, default=25, help="noise level only for DnCNN-S")
|
||||
parser.add_argument('--ckpt_prefix', type=str, default="dncnn_mindspore", help='ckpt name prefix')
|
||||
parser.add_argument('--epoch_num', type=int, default=50, help='epoch number')
|
||||
args = parser.parse_args()
|
||||
|
||||
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
|
||||
|
||||
if args.model_type == 'DnCNN-S':
|
||||
network = DnCNN(1, num_of_layers=17)
|
||||
elif args.model_type == 'DnCNN-3' or args.model_type == 'DnCNN-B':
|
||||
network = DnCNN(1, num_of_layers=20)
|
||||
else:
|
||||
print("wrong model type")
|
||||
exit()
|
||||
|
||||
ds_train = create_train_dataset(args.dataset_path, args.model_type, noise_level=args.noise_level, \
|
||||
batch_size=args.batch_size)
|
||||
opt = nn.AdamWeightDecay(network.trainable_params(), args.lr, weight_decay=args.weight_decay)
|
||||
loss_fun = BatchAverageMSELoss(args.batch_size)
|
||||
|
||||
model = Model(network, loss_fun, opt)
|
||||
|
||||
#training callbacks
|
||||
checkpoint_config = CheckpointConfig(save_checkpoint_steps=1000, keep_checkpoint_max=3)
|
||||
ckpoint_cb = ModelCheckpoint(prefix=args.ckpt_prefix, directory='./ckpt/', config=checkpoint_config)
|
||||
print_cb = Print_info()
|
||||
lr_cb = LearningRateScheduler(learning_rate_function)
|
||||
loss_monitor_cb = LossMonitor(per_print_times=100)
|
||||
|
||||
print(datetime.datetime.now(), " training starts")
|
||||
model.train(args.epoch_num, ds_train, callbacks=[lr_cb, ckpoint_cb, print_cb, loss_monitor_cb], \
|
||||
dataset_sink_mode=False)
|
||||
Loading…
Reference in New Issue