├── DRD-Net ├── Rain200H_test.py ├── Rain200L_test.py ├── generator.py ├── model.py ├── noise_model.py ├── rain800_test.py └── train.py └── README.md /DRD-Net/Rain200H_test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | import argparse 3 | import numpy as np 4 | from pathlib import Path 5 | import cv2 6 | import os 7 | from model import get_model 8 | from skimage.measure import compare_psnr 9 | 10 | 11 | 12 | def get_args(): 13 | parser = argparse.ArgumentParser(description="Test trained model", 14 | formatter_class=argparse.ArgumentDefaultsHelpFormatter) 15 | parser.add_argument("--image_dir_rain", type=str, default='dataset/200H/Rain_H/test/x_in', 16 | help="rain image dir") 17 | parser.add_argument("--image_dir_gt", type=str, default='dataset/200H/Rain_H/test/y_out', 18 | help="gt image dir") 19 | parser.add_argument("--model", type=str, default="the_end", 20 | help="model architecture ") 21 | parser.add_argument("--weight_file", type=str, default='./weights/rain200H/weights.106-0.002-28.16601.hdf5', 22 | help="trained weight file") 23 | parser.add_argument("--If_n", type=bool, default=True, 24 | help="If normalizing the image") 25 | parser.add_argument("--output_dir", type=str, default='./result', 26 | help="if set, save resulting images otherwise show result using imshow") 27 | args = parser.parse_args() 28 | return args 29 | 30 | 31 | def get_image(image): 32 | image = np.clip(image, 0, 255) 33 | return np.uint8(image) 34 | 35 | 36 | def main(): 37 | args = get_args() 38 | image_dir_rain = args.image_dir_rain 39 | image_dir_gt = args.image_dir_gt 40 | weight_file = args.weight_file 41 | if_n = args.If_n 42 | model = get_model(args.model) 43 | model.load_weights(weight_file) 44 | if args.output_dir: 45 | output_dir = Path(args.output_dir) 46 | output_dir.mkdir(parents=True, exist_ok=True) 47 | 48 | image_paths = list(Path(image_dir_rain).glob("*.*")) 49 | psnr_all = [] 50 | for image_path_rain in image_paths: 51 | name = str(image_path_rain).split("\\")[-1] 52 | image_rain = cv2.imread(str(image_path_rain)) 53 | image_gt = cv2.imread(str(os.path.join(image_dir_gt, name))) 54 | if if_n: 55 | image_rain = image_rain/255.0 56 | noise_image = image_rain 57 | pred = model.predict(np.expand_dims(noise_image, 0)) 58 | if if_n: 59 | denoised_image = get_image(pred[1][0]*[255]) 60 | else: 61 | denoised_image = get_image(pred[1][0]) 62 | psnr_all.append(compare_psnr(denoised_image, image_gt)) 63 | 64 | 65 | if args.output_dir: 66 | cv2.imwrite(str(output_dir.joinpath(name)), denoised_image) 67 | 68 | print(np.mean(psnr_all)) 69 | 70 | if __name__ == '__main__': 71 | main() 72 | -------------------------------------------------------------------------------- /DRD-Net/Rain200L_test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | import argparse 3 | import numpy as np 4 | from pathlib import Path 5 | import cv2 6 | import os 7 | from model import get_model 8 | from skimage.measure import compare_psnr 9 | 10 | 11 | 12 | def get_args(): 13 | parser = argparse.ArgumentParser(description="Test trained model", 14 | formatter_class=argparse.ArgumentDefaultsHelpFormatter) 15 | parser.add_argument("--image_dir_rain", type=str, default='dataset/200L/test/Rain100L', 16 | help="rain image dir") 17 | parser.add_argument("--image_dir_gt", type=str, default='dataset/200L/test/noRain100L', 18 | help="gt image dir") 19 | parser.add_argument("--model", type=str, default="the_end", 20 | help="model architecture ") 21 | parser.add_argument("--weight_file", type=str, default='./weights/rain200L/weights.120-15.616-37.28831.hdf5', 22 | help="trained weight file") 23 | parser.add_argument("--If_n", type=bool, default=False, 24 | help="If normalizing the image") 25 | parser.add_argument("--output_dir", type=str, default='./result', 26 | help="if set, save resulting images otherwise show result using imshow") 27 | args = parser.parse_args() 28 | return args 29 | 30 | 31 | def get_image(image): 32 | image = np.clip(image, 0, 255) 33 | return np.uint8(image) 34 | 35 | 36 | def main(): 37 | args = get_args() 38 | image_dir_rain = args.image_dir_rain 39 | image_dir_gt = args.image_dir_gt 40 | weight_file = args.weight_file 41 | if_n = args.If_n 42 | model = get_model(args.model) 43 | model.load_weights(weight_file) 44 | model.summary() 45 | if args.output_dir: 46 | output_dir = Path(args.output_dir) 47 | output_dir.mkdir(parents=True, exist_ok=True) 48 | 49 | image_paths = list(Path(image_dir_rain).glob("*.*")) 50 | psnr_all = [] 51 | for image_path_rain in image_paths: 52 | name = str(image_path_rain).split("\\")[-1] 53 | image_rain = cv2.imread(str(image_path_rain)) 54 | image_gt = cv2.imread(str(os.path.join(image_dir_gt, name))) 55 | if if_n: 56 | image_rain = image_rain/255.0 57 | noise_image = image_rain 58 | pred = model.predict(np.expand_dims(noise_image, 0)) 59 | if if_n: 60 | denoised_image = get_image(pred[1][0]*[255]) 61 | else: 62 | denoised_image = get_image(pred[1][0]) 63 | psnr_all.append(compare_psnr(denoised_image, image_gt)) 64 | 65 | 66 | if args.output_dir: 67 | cv2.imwrite(str(output_dir.joinpath(name)), denoised_image) 68 | 69 | print(np.mean(psnr_all)) 70 | 71 | if __name__ == '__main__': 72 | main() 73 | -------------------------------------------------------------------------------- /DRD-Net/generator.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | import random 3 | import numpy as np 4 | import cv2 5 | from keras.utils import Sequence 6 | 7 | 8 | class NoisyImageGenerator(Sequence): 9 | def __init__(self, image_dir_zao, image_dir_yuan, if_n=False, batch_size=32, image_size=64): 10 | self.image_zao_paths = list(Path(image_dir_zao).glob("*.*")) 11 | self.image_yuan_paths = list(Path(image_dir_yuan).glob("*.*")) 12 | self.image_num = len(self.image_zao_paths) 13 | self.if_n = if_n 14 | self.batch_size = batch_size 15 | self.image_size = image_size 16 | 17 | def __len__(self): 18 | return self.image_num // self.batch_size 19 | 20 | def __getitem__(self, idx): 21 | batch_size = self.batch_size 22 | image_size = self.image_size 23 | x = np.zeros((batch_size, image_size, image_size, 3), dtype=np.float) 24 | # rain = np.zeros((batch_size, image_size, image_size, 3), dtype=np.float) 25 | y = np.zeros((batch_size, image_size, image_size, 3), dtype=np.float) 26 | sample_id = 0 27 | 28 | while True: 29 | aa = random.randint(0, self.image_num - 1) 30 | image_zao_path = self.image_zao_paths[aa] 31 | image_yuan_path = self.image_yuan_paths[aa] 32 | image_zao = cv2.imread(str(image_zao_path)) 33 | image_yuan = cv2.imread(str(image_yuan_path)) 34 | if self.if_n: 35 | image_zao = image_zao / 255.0 36 | image_yuan = image_yuan / 255.0 37 | 38 | h, w, _ = image_zao.shape 39 | 40 | if h >= image_size and w >= image_size: 41 | h, w, _ = image_zao.shape 42 | i = np.random.randint(h - image_size + 1) 43 | j = np.random.randint(w - image_size + 1) 44 | zao_patch = image_zao[i:i + image_size, j:j + image_size] 45 | yuan_patch = image_yuan[i:i + image_size, j:j + image_size] 46 | # rain_patch = zao_patch - yuan_patch 47 | x[sample_id] = zao_patch.astype(np.float) 48 | y[sample_id] = yuan_patch.astype(np.float) 49 | # rain[sample_id] = rain_patch.astype(np.uint8) 50 | 51 | sample_id += 1 52 | 53 | if sample_id == batch_size: 54 | return [{'Rain_image': x}, {'subtract_1': y, 'add_36': y}] 55 | 56 | 57 | class ValGenerator(Sequence): 58 | def __init__(self, image_dir_zao, image_dir_yuan, if_n=False): 59 | image_zao_paths = list(Path(image_dir_zao).glob("*.*")) 60 | image_yuan_paths = list(Path(image_dir_yuan).glob("*.*")) 61 | self.image_num = len(image_zao_paths) 62 | self.if_n = if_n 63 | self.data = [] 64 | 65 | for i_image in range(self.image_num): 66 | x = cv2.imread(str(image_zao_paths[i_image])) 67 | y = cv2.imread(str(image_yuan_paths[i_image])) 68 | if self.if_n: 69 | x = x / 255.0 70 | y = y / 255.0 71 | 72 | self.data.append([{'Rain_image': np.expand_dims(x, axis=0)}, {'subtract_1': np.expand_dims(y, axis=0), 73 | 'add_36': np.expand_dims(y, axis=0)}]) 74 | 75 | def __len__(self): 76 | return self.image_num 77 | 78 | def __getitem__(self, idx): 79 | return self.data[idx] 80 | -------------------------------------------------------------------------------- /DRD-Net/model.py: -------------------------------------------------------------------------------- 1 | from keras.models import Model 2 | from keras.layers import Input, Add,Subtract, PReLU, Conv2DTranspose, \ 3 | Concatenate, MaxPooling2D, UpSampling2D, Dropout, concatenate, GlobalAveragePooling2D,\ 4 | Reshape, Dense, multiply 5 | from keras.layers.convolutional import Conv2D 6 | from keras.layers.normalization import BatchNormalization 7 | from keras.callbacks import Callback 8 | from keras import backend as K 9 | import tensorflow as tf 10 | 11 | 12 | class L0Loss: 13 | def __init__(self): 14 | self.gamma = K.variable(2.) 15 | 16 | def __call__(self): 17 | def calc_loss(y_true, y_pred): 18 | loss = K.pow(K.abs(y_true - y_pred) + 1e-8, self.gamma) 19 | return loss 20 | return calc_loss 21 | 22 | 23 | class UpdateAnnealingParameter(Callback): 24 | def __init__(self, gamma, nb_epochs, verbose=0): 25 | super(UpdateAnnealingParameter, self).__init__() 26 | self.gamma = gamma 27 | self.nb_epochs = nb_epochs 28 | self.verbose = verbose 29 | 30 | def on_epoch_begin(self, epoch, logs=None): 31 | new_gamma = 2.0 * (self.nb_epochs - epoch) / self.nb_epochs 32 | K.set_value(self.gamma, new_gamma) 33 | 34 | if self.verbose > 0: 35 | print('\nEpoch %05d: UpdateAnnealingParameter reducing gamma to %s.' % (epoch + 1, new_gamma)) 36 | 37 | 38 | def tf_log10(x): 39 | numerator = tf.log(x) 40 | denominator = tf.log(tf.constant(10, dtype=numerator.dtype)) 41 | return numerator / denominator 42 | 43 | 44 | def PSNR(y_true, y_pred): 45 | y_true_t = y_true *255.0 46 | max_pixel = 255.0 47 | y_pred_t = K.clip(y_pred*255.0, 0.0, 255.0) 48 | # y_pred = K.clip(y_pred, 0.0, 1.0) 49 | return 10.0 * tf_log10((max_pixel ** 2) / (K.mean(K.square(y_pred_t - y_true_t)))) 50 | 51 | 52 | def get_model(model_name="the_end"): 53 | 54 | if model_name == "the_end": 55 | return get_the_end_model() 56 | else: 57 | raise ValueError("model_name should be 'srresnet'or 'unet'") 58 | 59 | 60 | def get_the_end_model(input_channel_num=3, feature_dim=64, resunit_num=16): 61 | 62 | def _back_net(inputs): 63 | def _residual_block(inputs): 64 | x = _empty_block(inputs) 65 | x = BatchNormalization()(x) 66 | x = PReLU(shared_axes=[1, 2])(x) 67 | x = _empty_block(x) 68 | x = BatchNormalization()(x) 69 | m = Add()([x, inputs]) 70 | return m 71 | 72 | def _empty_block(inputs): 73 | x1 = Conv2D(feature_dim, (3, 3), padding="same", kernel_initializer="he_normal")(inputs) 74 | x2 = Conv2D(feature_dim, (3, 3), dilation_rate=3, padding="same", kernel_initializer="he_normal")(inputs) 75 | x3 = Conv2D(feature_dim, (3, 3), dilation_rate=5, padding="same", kernel_initializer="he_normal")(inputs) 76 | x = concatenate([x1, x2, x3], axis=-1) 77 | x_out = Conv2D(feature_dim, (1, 1), padding="same", kernel_initializer="he_normal")(x) 78 | return x_out 79 | 80 | x = Conv2D(feature_dim, (3, 3), padding="same", kernel_initializer="he_normal")(inputs) 81 | x = PReLU(shared_axes=[1, 2])(x) 82 | x0 = x 83 | 84 | for i in range(resunit_num): 85 | x = _residual_block(x) 86 | 87 | x = Conv2D(feature_dim, (3, 3), padding="same", kernel_initializer="he_normal")(x) 88 | x = BatchNormalization()(x) 89 | x = Add()([x, x0]) 90 | x = Conv2D(input_channel_num, (3, 3), padding="same", kernel_initializer="he_normal")(x) 91 | return x 92 | 93 | def _rain_net(inputs): 94 | def _residual_block(inputs, number): 95 | x = Conv2D(feature_dim, (3, 3), padding="same", kernel_initializer="he_normal")(inputs) 96 | x = BatchNormalization()(x) 97 | x = PReLU(shared_axes=[1, 2])(x) 98 | x = Conv2D(feature_dim, (3, 3), padding="same", kernel_initializer="he_normal")(x) 99 | x = BatchNormalization()(x) 100 | filters = 64 101 | se_shape = (1, 1, filters) 102 | se = GlobalAveragePooling2D()(x) 103 | se = Reshape(se_shape)(se) 104 | se = Dense(number, activation="relu", 105 | kernel_initializer="he_normal",use_bias=False)(se) 106 | se = Dense(filters, activation="hard_sigmoid", 107 | kernel_initializer="he_normal", use_bias=False)(se) 108 | x = multiply([x, se]) 109 | m = Add()([x, inputs]) 110 | return m 111 | 112 | x = Conv2D(feature_dim, (3, 3), padding="same", kernel_initializer="he_normal")(inputs) 113 | x = PReLU(shared_axes=[1, 2])(x) 114 | x0 = x 115 | 116 | for i in range(resunit_num): 117 | x = _residual_block(x, 4) 118 | 119 | x = Conv2D(feature_dim, (3, 3), padding="same", kernel_initializer="he_normal")(x) 120 | x = BatchNormalization()(x) 121 | x = Add()([x, x0]) 122 | x = Conv2D(input_channel_num, (3, 3), padding="same", kernel_initializer="he_normal")(x) 123 | return x 124 | 125 | inputs = Input(shape=(None, None, input_channel_num), name='Rain_image') 126 | Rain = _rain_net(inputs) 127 | 128 | out1 = Subtract()([inputs, Rain]) 129 | Back_in = Add()([out1, inputs]) 130 | Back = _back_net(Back_in) 131 | out = Add()([out1, Back]) 132 | 133 | model = Model(inputs=inputs, outputs=[out1, out]) 134 | return model 135 | 136 | 137 | def main(): 138 | # model = get_model() 139 | model = get_model("unet") 140 | model.summary() 141 | 142 | 143 | if __name__ == '__main__': 144 | main() 145 | -------------------------------------------------------------------------------- /DRD-Net/noise_model.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import string 3 | import random 4 | import numpy as np 5 | import cv2 6 | import skimage 7 | 8 | 9 | 10 | def get_noise_model(noise_type="gaussian,0,50"): 11 | tokens = noise_type.split(sep=",") 12 | 13 | if tokens[0] == "gaussian": 14 | min_stddev = int(tokens[1]) 15 | max_stddev = int(tokens[2]) 16 | 17 | def gaussian_noise(img): 18 | noise_img = img.astype(np.float) 19 | stddev = np.random.uniform(min_stddev, max_stddev) 20 | noise = np.random.randn(*img.shape) * stddev 21 | noise_img += noise 22 | noise_img = np.clip(noise_img, 0, 255).astype(np.uint8) 23 | return noise_img 24 | return gaussian_noise 25 | elif tokens[0] == "clean": 26 | return lambda img: img 27 | elif tokens[0] == "text": 28 | min_occupancy = int(tokens[1]) 29 | max_occupancy = int(tokens[2]) 30 | 31 | def add_text(img): 32 | img = img.copy() 33 | h, w, _ = img.shape 34 | font = cv2.FONT_HERSHEY_SIMPLEX 35 | img_for_cnt = np.zeros((h, w), np.uint8) 36 | occupancy = np.random.uniform(min_occupancy, max_occupancy) 37 | 38 | while True: 39 | n = random.randint(5, 10) 40 | random_str = ''.join([random.choice(string.ascii_letters + string.digits) for i in range(n)]) 41 | font_scale = np.random.uniform(0.5, 1) 42 | thickness = random.randint(1, 3) 43 | (fw, fh), baseline = cv2.getTextSize(random_str, font, font_scale, thickness) 44 | x = random.randint(0, max(0, w - 1 - fw)) 45 | y = random.randint(fh, h - 1 - baseline) 46 | color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) 47 | cv2.putText(img, random_str, (x, y), font, font_scale, color, thickness) 48 | cv2.putText(img_for_cnt, random_str, (x, y), font, font_scale, 255, thickness) 49 | 50 | if (img_for_cnt > 0).sum() > h * w * occupancy / 100: 51 | break 52 | return img 53 | return add_text 54 | elif tokens[0] == "impulse": 55 | min_occupancy = int(tokens[1]) 56 | max_occupancy = int(tokens[2]) 57 | 58 | # def add_impulse_noise(img): 59 | # occupancy = np.random.uniform(min_occupancy, max_occupancy) 60 | # mask = np.random.binomial(size=img.shape, n=1, p=occupancy / 100) 61 | # noise = np.random.randint(256, size=img.shape) 62 | # img = img * (1 - mask) + noise * mask 63 | # return img.astype(np.uint8) 64 | def add_impulse_noise(img): 65 | if min_occupancy == 3: 66 | return img.astype(np.uint8) 67 | elif min_occupancy == 1: 68 | return img.astype(np.uint8) 69 | return add_impulse_noise 70 | else: 71 | raise ValueError("noise_type should be 'gaussian', 'clean', 'text', or 'impulse'") 72 | 73 | 74 | def get_args(): 75 | parser = argparse.ArgumentParser(description="test noise model", 76 | formatter_class=argparse.ArgumentDefaultsHelpFormatter) 77 | parser.add_argument("--image_size", type=int, default=256, 78 | help="training patch size") 79 | parser.add_argument("--noise_model", type=str, default="gaussian,0,50", 80 | help="noise model to be tested") 81 | args = parser.parse_args() 82 | return args 83 | 84 | 85 | def main(): 86 | args = get_args() 87 | image_size = args.image_size 88 | noise_model = get_noise_model(args.noise_model) 89 | 90 | while True: 91 | image = np.ones((image_size, image_size, 3), dtype=np.uint8) * 128 92 | noisy_image = noise_model(image) 93 | cv2.imshow("noise image", noisy_image) 94 | key = cv2.waitKey(-1) 95 | 96 | # "q": quit 97 | if key == 113: 98 | return 0 99 | 100 | 101 | if __name__ == '__main__': 102 | main() 103 | -------------------------------------------------------------------------------- /DRD-Net/rain800_test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | import argparse 3 | import numpy as np 4 | from pathlib import Path 5 | import cv2 6 | import os 7 | from model import get_model 8 | from skimage.measure import compare_psnr 9 | 10 | 11 | 12 | def get_args(): 13 | parser = argparse.ArgumentParser(description="Test trained model", 14 | formatter_class=argparse.ArgumentDefaultsHelpFormatter) 15 | parser.add_argument("--image_dir_rain", type=str, default='dataset/rain800/test/norain', 16 | help="rain image dir") 17 | parser.add_argument("--image_dir_gt", type=str, default='dataset/rain800/test/rain', 18 | help="gt image dir") 19 | parser.add_argument("--model", type=str, default="the_end", 20 | help="model architecture ") 21 | parser.add_argument("--weight_file", type=str, default='./weights/rain800/weights.107-201.542-26.45626.hdf5', 22 | help="trained weight file") 23 | parser.add_argument("--If_n", type=bool, default=False, 24 | help="If normalizing the image") 25 | parser.add_argument("--output_dir", type=str, default='./result', 26 | help="if set, save resulting images otherwise show result using imshow") 27 | args = parser.parse_args() 28 | return args 29 | 30 | 31 | def get_image(image): 32 | image = np.clip(image, 0, 255) 33 | return np.uint8(image) 34 | 35 | 36 | def main(): 37 | args = get_args() 38 | image_dir_rain = args.image_dir_rain 39 | image_dir_gt = args.image_dir_gt 40 | weight_file = args.weight_file 41 | if_n = args.If_n 42 | model = get_model(args.model) 43 | model.load_weights(weight_file) 44 | if args.output_dir: 45 | output_dir = Path(args.output_dir) 46 | output_dir.mkdir(parents=True, exist_ok=True) 47 | 48 | image_paths = list(Path(image_dir_rain).glob("*.*")) 49 | psnr_all = [] 50 | for image_path_rain in image_paths: 51 | name = str(image_path_rain).split("\\")[-1] 52 | image_rain = cv2.imread(str(image_path_rain)) 53 | image_gt = cv2.imread(str(os.path.join(image_dir_gt, name))) 54 | if if_n: 55 | image_rain = image_rain/255.0 56 | noise_image = image_rain 57 | pred = model.predict(np.expand_dims(noise_image, 0)) 58 | if if_n: 59 | denoised_image = get_image(pred[1][0]*[255]) 60 | else: 61 | denoised_image = get_image(pred[1][0]) 62 | psnr_all.append(compare_psnr(denoised_image, image_gt)) 63 | 64 | 65 | if args.output_dir: 66 | cv2.imwrite(str(output_dir.joinpath(name)), denoised_image) 67 | 68 | print(np.mean(psnr_all)) 69 | 70 | if __name__ == '__main__': 71 | main() 72 | -------------------------------------------------------------------------------- /DRD-Net/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import numpy as np 3 | from pathlib import Path 4 | from keras.callbacks import LearningRateScheduler, ModelCheckpoint,TensorBoard, ReduceLROnPlateau 5 | from keras.optimizers import Adam 6 | from model import get_model, PSNR, L0Loss, UpdateAnnealingParameter 7 | from generator import NoisyImageGenerator, ValGenerator 8 | 9 | 10 | class Schedule: 11 | def __init__(self, nb_epochs, initial_lr): 12 | self.epochs = nb_epochs 13 | self.initial_lr = initial_lr 14 | 15 | def __call__(self, epoch_idx): 16 | if epoch_idx < 6: 17 | return 0.01 18 | elif epoch_idx < 15: 19 | return 0.005 20 | elif epoch_idx < 30: 21 | return 0.0025 22 | elif epoch_idx < 45: 23 | return 0.00125 24 | elif epoch_idx < 60: 25 | return 0.001 26 | elif epoch_idx < 75: 27 | return 0.0005 28 | elif epoch_idx < 90: 29 | return 0.000125 30 | return 0.00001 31 | 32 | 33 | def get_args(): 34 | parser = argparse.ArgumentParser(description="train derain model", 35 | formatter_class=argparse.ArgumentDefaultsHelpFormatter) 36 | parser.add_argument("--image_dir_noise", type=str, default='./dataset/Rain_H/train/x_in', 37 | help="train image dir") 38 | parser.add_argument("--image_dir_original", type=str, default='./dataset/Rain_H/train/y_out', 39 | help="train image dir") 40 | parser.add_argument("--test_dir_noise", type=str, default='./dataset/Rain100H/rain', 41 | help="test image dir") 42 | parser.add_argument("--test_dir_original", type=str, default='./dataset/Rain100H/norain', 43 | help="test image dir") 44 | parser.add_argument("--If_n", type=bool, default=True, 45 | help="If normalizing the image") 46 | parser.add_argument("--image_size", type=int, default=128, 47 | help="training patch size") 48 | parser.add_argument("--batch_size", type=int, default=4, 49 | help="batch size") 50 | parser.add_argument("--nb_epochs", type=int, default=120, 51 | help="number of epochs") 52 | parser.add_argument("--lr", type=float, default=0.1, 53 | help="learning rate") 54 | parser.add_argument("--steps", type=int, default=1000, 55 | help="steps per epoch") 56 | parser.add_argument("--loss", type=str, default="mse", 57 | help="loss; mse', 'mae', or 'l0' is expected") 58 | parser.add_argument("--output_path", type=str, default="impulse_clean", 59 | help="checkpoint dir") 60 | parser.add_argument("--model", type=str, default="the_end", 61 | help="model architecture ('Similarity')") 62 | args = parser.parse_args() 63 | 64 | return args 65 | 66 | 67 | def main(): 68 | args = get_args() 69 | image_dir_noise = args.image_dir_noise 70 | image_dir_original = args.image_dir_original 71 | test_dir_noise = args.test_dir_noise 72 | test_dir_original = args.test_dir_original 73 | image_size = args.image_size 74 | batch_size = args.batch_size 75 | nb_epochs = args.nb_epochs 76 | if_n = args.If_n 77 | lr = args.lr 78 | steps = args.steps 79 | loss_type = args.loss 80 | output_path = Path(__file__).resolve().parent.joinpath(args.output_path) 81 | model = get_model(args.model) 82 | opt = Adam(lr=lr) 83 | callbacks = [] 84 | 85 | if loss_type == "l0": 86 | l0 = L0Loss() 87 | callbacks.append(UpdateAnnealingParameter(l0.gamma, nb_epochs, verbose=1)) 88 | loss_type = l0() 89 | 90 | # model.compile(optimizer=opt, loss=loss_type, metrics=[PSNR]) 91 | model.compile(optimizer=opt, loss={"subtract_1": "mse", "add_36": "mse"}, 92 | loss_weights={'subtract_1': 0.1, 'add_36': 1}, metrics=[PSNR]) 93 | model.summary() 94 | generator = NoisyImageGenerator(image_dir_noise, image_dir_original, if_n=if_n, batch_size=batch_size, 95 | image_size=image_size) 96 | val_generator = ValGenerator(test_dir_noise, test_dir_original, if_n=if_n) 97 | output_path.mkdir(parents=True, exist_ok=True) 98 | # callbacks.append(ReduceLROnPlateau(monitor='val_add_35_loss', factor=0.5, patience=5, verbose=1, mode='min', 99 | # cooldown=0, min_lr=0.000000001)) 100 | callbacks.append(LearningRateScheduler(schedule=Schedule(nb_epochs, lr))) 101 | callbacks.append(TensorBoard(log_dir='./log', histogram_freq=0, batch_size=batch_size, write_images=True)) 102 | callbacks.append(ModelCheckpoint(str(output_path) + "/weights.{epoch:03d}-{val_add_36_loss:.3f}-{val_add_36_PSNR:.5f}.hdf5", 103 | monitor="val_add_36_PSNR", 104 | verbose=1, 105 | mode="max", 106 | save_best_only=True)) 107 | 108 | hist = model.fit_generator(generator=generator, 109 | steps_per_epoch=steps, 110 | epochs=nb_epochs, 111 | validation_data=val_generator, 112 | verbose=1, 113 | callbacks=callbacks) 114 | 115 | np.savez(str(output_path.joinpath("history.npz")), history=hist.history) 116 | 117 | 118 | if __name__ == '__main__': 119 | main() 120 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # DRD-Net 2 | This website shares the codes of the "Detail-recovery Image Deraining via Context Aggregation Networks",CVPR 2020. 3 | 4 | Prerequisites: 5 | tensorflow == 1.10.0 6 | keras == 2.2.4 7 | python == 3.6 8 | CUDA ==10.0 9 | 10 | For train the Rain200H, please run: 11 | python train.py --image_dir_noise you rain data --image_dir_original you gt data --test_dir_noise you test rain data --test_dir_original you test gt data --If_n True 12 | 13 | For train the Rain200L, please run: 14 | python train.py --image_dir_noise you rain data --image_dir_original you gt data --test_dir_noise you test rain data --test_dir_original you test gt data --If_n False 15 | 16 | For train the Rain800, please run: 17 | python train.py --image_dir_noise you rain data --image_dir_original you gt data --test_dir_noise you test rain data --test_dir_original you test gt data --If_n False 18 | 19 | The dataset "Rain200H" and "Rain200L" you can download here: 20 | https://www.icst.pku.edu.cn/struct/Projects/joint_rain_removal.html 21 | 22 | The dataset "Rain800" you can download here: 23 | https://drive.google.com/drive/folders/0Bw2e6Q0nQQvGbi1xV1Yxd09rY2s 24 | --------------------------------------------------------------------------------