[ { "path": "README.md", "content": "# Learnbale_Bandpass_Filter\nImage Demoireing with Learnable Bandpass Filters, CVPR2020\n\nOur extension work is accepted by IEEE TPAMI.\nThe journal paper will come soon.\n\nIf you find this work is helpful, please cite:\n\n@article{zheng2021learn,\n \n title={Learning Frequency Domain Priors for Image Demoireing},\n \n author = {Bolun, Zheng and Shanxin, Yuan and Chenggang, Yan and Xiang, Tian and Jiyong, Zhang and Yaoqi, Sun and Lin, Liu and Ales, Leonardis and Gregory, Slabaugh},\n \n journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},\n \n year={2021}\n}\n\n@inProceedings{zheng2020, \n\nauthor={B. Zheng and S. Yuan and G. Slabaugh and A. Leonardis}, \n\nbooktitle={IEEE Conference on Computer Vision and Pattern Recongnition}, \n\ntitle={Image Demoireing with Learnable Bandpass Filters}, \n\nyear={2020}, \n}\n\n@article{zheng2019implicit,\n \n title={Implicit dual-domain convolutional network for robust color image compression artifact reduction},\n \n author={Zheng, Bolun and Chen, Yaowu and Tian, Xiang and Zhou, Fan and Liu, Xuesong},\n \n journal={IEEE Transactions on Circuits and Systems for Video Technology},\n \n volume={30},\n \n number={11},\n \n pages={3982--3994},\n \n year={2020},\n \n publisher={IEEE}\n}\n\nYou can now get this paper at Arxiv preprint: https://arxiv.org/abs/2004.00406\n## Run the code\nThis project requires:\n* Tensorflow >1.10\n* Keras > 2.0\n* opencv > 2.0\n* skImage\n\nYou can get the weight file for AIM2019 via: \nhttps://1drv.ms/u/s!ArU0YIIFiFuHilwyuwHZjSpvPUBz?e=iZ70Ga \nor via Baidu Disk： \nhttps://pan.baidu.com/s/1wsJYyYbQO-ETL5Jq4fN6hw code：jiae \n\nYou can get AIM2019 LCDMoire2019 dataset via:\nvalidation: \nMoire: https://data.vision.ee.ethz.ch/timofter/AIM19demoire/ValidationMoire.zip \nClean: https://data.vision.ee.ethz.ch/timofter/AIM19demoire/ValidationClear.zip \n\ntesting: \nhttps://data.vision.ee.ethz.ch/timofter/AIM19demoire/TestingMoire.zip\n\n\nThen, \n1. edit the 'main_multiscale.py' by:\nreplacing the 'test_path', 'valid_gt_path', 'valid_ns_path' and 'weight_path' with your own settings. \n\n2. make the dirs 'testing_result' and 'validation_result' at current path. \n\n3. python main_multiscale.py. \n\n" }, { "path": "core_layers.py", "content": "import tensorflow as tf\r\nfrom keras import backend as k\r\nfrom keras.utils import conv_utils\r\nimport numpy as np\r\nfrom keras import layers, models, activations, initializers, constraints\r\nfrom math import cos, pi, sqrt\r\nfrom keras.regularizers import l2\r\n\r\nclass Space2Depth(layers.Layer):\r\n def __init__(self, scale, **kwargs):\r\n super(Space2Depth, self).__init__(**kwargs)\r\n self.scale = scale\r\n\r\n def call(self, inputs, **kwargs):\r\n return tf.space_to_depth(inputs, self.scale)\r\n\r\n def compute_output_shape(self, input_shape):\r\n if input_shape[1] != None and input_shape[2] != None:\r\n return (None, int(input_shape[1]/self.scale), int(input_shape[2]/self.scale), input_shape[3]*self.scale**2)\r\n else:\r\n return (None, None, None, input_shape[3]*self.scale**2)\r\n\r\nclass Depth2Space(layers.Layer):\r\n def __init__(self, scale, **kwargs):\r\n super(Depth2Space, self).__init__(**kwargs)\r\n self.scale = scale\r\n def call(self, inputs, **kwargs):\r\n return tf.depth_to_space(inputs, self.scale)\r\n\r\n def compute_output_shape(self, input_shape):\r\n if input_shape[1] != None and input_shape[2] != None:\r\n return (None, input_shape[1]*self.scale, input_shape[2]*self.scale, int(input_shape[3]/self.scale**2))\r\n else:\r\n return (None, None, None, int(input_shape[3]/self.scale**2))\r\n\r\nclass adaptive_implicit_trans(layers.Layer):\r\n def __init__(self, **kwargs):\r\n super(adaptive_implicit_trans, self).__init__(**kwargs)\r\n\r\n def build(self, input_shape):\r\n conv_shape = (1,1,64,64)\r\n self.it_weights = self.add_weight(\r\n shape = (1,1,64,1),\r\n initializer = initializers.get('ones'),\r\n constraint = constraints.NonNeg(),\r\n name = 'ait_conv')\r\n kernel = np.zeros(conv_shape)\r\n r1 = sqrt(1.0/8)\r\n r2 = sqrt(2.0/8)\r\n for i in range(8):\r\n _u = 2*i+1\r\n for j in range(8):\r\n _v = 2*j+1\r\n index = i*8+j\r\n for u in range(8):\r\n for v in range(8):\r\n index2 = u*8+v\r\n t = cos(_u*u*pi/16)*cos(_v*v*pi/16)\r\n t = t*r1 if u==0 else t*r2\r\n t = t*r1 if v==0 else t*r2\r\n kernel[0,0,index2,index] = t\r\n self.kernel = k.variable(value = kernel, dtype = 'float32')\r\n\r\n def call(self, inputs):\r\n #it_weights = k.softmax(self.it_weights)\r\n #self.kernel = self.kernel*it_weights\r\n self.kernel = self.kernel*self.it_weights\r\n y = k.conv2d(inputs,\r\n self.kernel,\r\n padding = 'same',\r\n data_format='channels_last')\r\n return y\r\n\r\n def compute_output_shape(self, input_shape):\r\n return input_shape\r\n\r\nclass ScaleLayer(layers.Layer):\r\n def __init__(self, s, **kwargs):\r\n self.s = s\r\n super(ScaleLayer, self).__init__(**kwargs)\r\n\r\n def build(self, input_shape):\r\n self.kernel = self.add_weight(\r\n shape = (1,),\r\n name = 'scale',\r\n initializer=initializers.Constant(value=self.s))\r\n def call(self, inputs):\r\n return inputs*self.kernel\r\n\r\n def compute_output_shape(self, input_shape):\r\n return input_shape\r\n" }, { "path": "main_multiscale.py", "content": "import tensorflow as tf\r\nimport os, cv2, random, keras\r\nfrom skimage.measure import compare_ssim\r\nimport numpy as np\r\nfrom model import *\r\nfrom keras import optimizers\r\nfrom keras.utils import multi_gpu_model\r\nfrom math import log10\r\nfrom utils import *\r\nfrom datetime import datetime\r\nimport time\r\nfrom PIL import Image\r\n\r\n#the image dir of testing input\r\ntest_path = 'F:\\\\dataset\\\\AIM2019\\\\Testing\\\\'\r\n#the image dir of validation groundtruth\r\nvalid_gt_path = 'Validation\\\\clear\\\\'\r\n#the image dir of validation input\r\nvalid_ns_path = 'Validation\\\\moire\\\\'\r\n#weight file path\r\nweight_path = 'model/MBCNN_weights.h5'\r\n\r\nmulti_input = False\r\nmulti_output = True\r\nmodel = MBCNN(64,multi_output) #MBCNN-light: model = MBCNN(32,multi_output)\r\n\r\n\r\n# Validation or Testing\r\ndef validate_ssim(model, gt_list, ns_list, name_list, multi_output=False):\r\n print (\"validating... \",datetime.now().strftime('%H:%M:%S'))\r\n psnr = 0\r\n ssim = 0\r\n count = 0\r\n for i in range(len(gt_list)):\t\r\n count += 1\r\n gt = gt_list[i]\r\n ns = ns_list[i]\t\r\n dn = model.predict(ns)#[-1]\r\n if multi_output:\r\n dn = dn[-1]\r\n _psnr = 10*log10(1/np.mean((dn-gt)**2))\r\n _ssim = compare_ssim(dn[0],gt[0],multichannel=True)\r\n psnr += _psnr\r\n ssim += _ssim\r\n _gt = dn[0]\r\n _gt[_gt>1] = 1\r\n _gt[_gt<0] = 0\r\n _gt = _gt*255.0\r\n _gt = np.round(_gt).astype(np.uint8)\r\n cv2.imwrite('validation_result/'+name_list[i],_gt)\r\n\r\n print (np.round(psnr/count,3),np.round(ssim/count,4))\r\n return psnr/count\r\n\r\ndef test(model,multi_output=False):\r\n psnr = 0\r\n count = 0\r\n \r\n file_list = os.listdir(test_path)\r\n file_list = list_filter(file_list,'.png')\r\n for f in file_list:\r\n ns = cv2.imread(test_path+f)\r\n ns = ns.astype(np.float32)/255.0\r\n ns = ns.reshape((1,)+ns.shape)\r\n start = time.clock()\r\n _gt = model.predict(ns)\r\n if multi_output:\r\n _gt = _gt[-1]\r\n end = time.clock()\r\n print(f, end-start)\r\n _gt = _gt[0]\r\n _gt[_gt>1] = 1\r\n _gt[_gt<0] = 0\r\n _gt = _gt*255.0\r\n _gt = np.round(_gt).astype(np.uint8)\r\n cv2.imwrite('testing_result/'+f,_gt)\r\n\r\n#Generating validation datas\r\ndef generate_validation(valid_list, multi_input, mode='sub'):\r\n valid_gt_list = []\r\n valid_ns_list = []\r\n name_list = []\r\n _width = 128\r\n for f in valid_list:\r\n _name = os.path.splitext(f)[0]\r\n gt = cv2.imread(valid_gt_path+f)\r\n ns = cv2.imread(valid_ns_path+f)\r\n gt = gt.astype(np.float32)/255.0\r\n ns = ns.astype(np.float32)/255.0\r\n name_list.append(f)\r\n if multi_input:\r\n ns_x2 = cv2.imread(ns_path+'X2\\\\'+_name+'.png')\r\n ns_x4 = cv2.imread(ns_path+'X4\\\\'+_name+'.png')\r\n ns_x8 = cv2.imread(ns_path+'X8\\\\'+_name+'.png')\r\n ns_x2 = ns_x2.astype(np.float32)/255.0\r\n ns_x4 = ns_x4.astype(np.float32)/255.0\r\n ns_x8 = ns_x8.astype(np.float32)/255.0\r\n \r\n if mode == 'sub':\r\n for i in range(0,1024,_width):\r\n for j in range(0,1024,_width):\r\n _gt = gt[i:i+_width,j:j+_width]\r\n _ns = ns[i:i+_width,j:j+_width]\r\n _gt = _gt.reshape((1,)+_gt.shape)\r\n _ns = _ns.reshape((1,)+_ns.shape)\r\n valid_gt_list.append(_gt)\r\n valid_ns_list.append(_ns)\r\n if mode == 'full':\r\n gt = gt.reshape((1,)+gt.shape)\r\n ns = ns.reshape((1,)+ns.shape)\r\n valid_gt_list.append(gt)\r\n if multi_input:\r\n ns_x2 = ns_x2.reshape((1,)+ns_x2.shape)\r\n ns_x4 = ns_x4.reshape((1,)+ns_x4.shape)\r\n ns_x8 = ns_x8.reshape((1,)+ns_x8.shape)\r\n valid_ns_list.append([ns,ns_x2,ns_x4,ns_x8])\r\n else:\r\n valid_ns_list.append(ns)\r\n return valid_gt_list, valid_ns_list, name_list \r\n\r\n#Training\r\nmulti_gpu = False\r\nif multi_gpu:\t\r\n\tprint ('use multi gpu mode!')\r\n\tos.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0,1\"\r\nelse:\r\n\tos.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\r\nkeras.backend.tensorflow_backend.set_session(get_session())\r\n\r\nmodel.summary()\r\n#exit(0)\r\n\r\nvalid_list = os.listdir(valid_gt_path)\r\nvalid_list = list_filter(valid_list, '.png')\r\n\r\nvalid_gt_list, valid_ns_list, name_list = generate_validation(valid_list, multi_input, 'full')\r\nmodel.load_weights(weight_path, by_name = True)\r\n#output validation results\r\nmin_loss = validate_ssim(model, valid_gt_list, valid_ns_list, name_list, multi_output)\r\n#output testing results\r\ntest(model, multi_output)\r\nexit(0)\r\n" }, { "path": "model.py", "content": "from keras import layers\r\nfrom keras.models import Model\r\nfrom keras import backend as K\r\nfrom core_layers import *\r\n\r\ndef conv_relu(x, filters, kernel, padding='same', use_bias = True, dilation_rate=1, strides=(1,1)):\r\n if dilation_rate == 0:\r\n y = layers.Conv2D(filters,1,padding=padding,use_bias=use_bias,\r\n activation='relu')(x)\r\n else:\r\n y = layers.Conv2D(filters,kernel,padding=padding,use_bias=use_bias,\r\n dilation_rate=dilation_rate,\r\n strides=strides,\r\n activation='relu')(x)\r\n return y\r\n\r\ndef conv(x, filters, kernel, padding='same', use_bias=True, dilation_rate=1, strides = (1,1)):\r\n y = layers.Conv2D(filters,kernel,padding=padding,use_bias=use_bias,\r\n dilation_rate=dilation_rate, strides=strides)(x)\r\n return y\r\n\r\ndef conv_bn_relu(x, filters, kernel, padding='same', use_bias = True, dilation_rate=1):\r\n y = layers.Conv2D(filters,kernel,padding=padding,use_bias=use_bias,\r\n dilation_rate=dilation_rate)(x)\r\n y = layers.BatchNormalization(axis=-1)(y)\r\n y = layers.Activation('relu')(y)\r\n return y\r\n\r\ndef conv_prelu(x, filters, kernel, padding='same', use_bias=False, dilation_rate=1, strides = (1,1)):\r\n y = layers.Conv2D(filters,kernel,padding=padding,use_bias=use_bias,\r\n dilation_rate=dilation_rate, strides=strides)(x)\r\n y = layers.advanced_activations.PReLU()(y)\r\n return y\r\n\r\ndef MBCNN(nFilters, multi=True):\r\n conv_func = conv_relu\r\n def pre_block(x, d_list, enbale = True):\r\n t = x\r\n for i in range(len(d_list)):\r\n _t = conv_func(t, nFilters, 3, dilation_rate=d_list[i])\r\n t = layers.Concatenate(axis=-1)([_t,t])\r\n t = conv(t, 64, 3)\r\n t = adaptive_implicit_trans()(t)\r\n t = conv(t,nFilters*2,1)\r\n t = ScaleLayer(s=0.1)(t)\r\n if not enbale:\r\n t = layers.Lambda(lambda x: x*0)(t)\r\n t = layers.Add()([x,t])\r\n return t\r\n\r\n def pos_block(x, d_list):\r\n t = x\r\n for i in range(len(d_list)):\r\n _t = conv_func(t, nFilters, 3, dilation_rate=d_list[i])\r\n t = layers.Concatenate(axis=-1)([_t,t])\r\n t = conv_func(t, nFilters*2, 1)\r\n return t\r\n\r\n def global_block(x):\r\n t = layers.ZeroPadding2D(padding=(1,1))(x)\r\n t = conv_func(t, nFilters*4, 3, strides=(2,2))\r\n t = layers.GlobalAveragePooling2D()(t)\r\n t = layers.Dense(nFilters*16,activation='relu')(t)\r\n t = layers.Dense(nFilters*8, activation='relu')(t)\r\n t = layers.Dense(nFilters*4)(t)\r\n _t = conv_func(x, nFilters*4, 1)\r\n _t = layers.Multiply()([_t,t])\r\n _t = conv_func(_t, nFilters*2, 1)\r\n return _t\r\n\r\n output_list = []\r\n d_list_a = (1,2,3,2,1)\r\n d_list_b = (1,2,3,2,1)\r\n d_list_c = (1,2,2,2,1)\r\n x = layers.Input(shape=(None, None, 3)) #16m*16m\r\n _x = Space2Depth(scale=2)(x)\r\n t1 = conv_func(_x,nFilters*2,3, padding='same') #8m*8m\r\n t1 = pre_block(t1, d_list_a, True)\r\n t2 = layers.ZeroPadding2D(padding=(1,1))(t1)\r\n t2 = conv_func(t2,nFilters*2,3, padding='valid',strides=(2,2)) #4m*4m\r\n t2 = pre_block(t2, d_list_b,True)\r\n t3 = layers.ZeroPadding2D(padding=(1,1))(t2)\r\n t3 = conv_func(t3,nFilters*2,3, padding='valid',strides=(2,2)) #2m*2m\r\n t3 = pre_block(t3,d_list_c, True)\r\n t3 = global_block(t3)\r\n t3 = pos_block(t3, d_list_c)\r\n t3_out = conv(t3, 12, 3)\r\n t3_out = Depth2Space(scale=2)(t3_out) #4m*4m\r\n output_list.append(t3_out)\r\n _t2 = layers.Concatenate()([t3_out,t2])\r\n _t2 = conv_func(_t2, nFilters*2, 1)\r\n _t2 = global_block(_t2)\r\n _t2 = pre_block(_t2, d_list_b,True)\r\n _t2 = global_block(_t2)\r\n _t2 = pos_block(_t2, d_list_b)\r\n t2_out = conv(_t2, 12, 3)\r\n t2_out = Depth2Space(scale=2)(t2_out) #8m*8m\r\n output_list.append(t2_out)\r\n _t1 = layers.Concatenate()([t1, t2_out])\r\n _t1 = conv_func(_t1, nFilters*2, 1)\r\n _t1 = global_block(_t1)\r\n _t1 = pre_block(_t1, d_list_a, True)\r\n _t1 = global_block(_t1)\r\n _t1 = pos_block(_t1, d_list_a)\r\n _t1 = conv(_t1,12,3)\r\n y = Depth2Space(scale=2)(_t1) #16m*16m\r\n output_list.append(y)\r\n if multi != True:\r\n return models.Model(x,y)\r\n else:\r\n return models.Model(x,output_list)\r\n" }, { "path": "utils.py", "content": "import numpy as np\r\nimport math, os\r\nfrom keras import optimizers, backend\r\nimport tensorflow as tf\r\nimport cv2\r\n\r\ndef get_Y(x):\r\n\tr = x[:,:,0]\r\n\tg = x[:,:,1]\r\n\tb = x[:,:,2]\r\n\ty = 0.257*r + 0.504*g + 0.098*b + 0.0627\r\n\treturn y\r\n\r\ndef calc_PSNR(x,y):\r\n\tmse = np.mean(np.square(x-y))\r\n\tpsnr = 10*math.log10(1/mse)\r\n\treturn psnr\r\n\r\ndef get_session():\r\n config = tf.ConfigProto()\r\n config.gpu_options.allow_growth = True\r\n return tf.Session(config = config)\r\n\r\ndef list_filter(file_list, tail):\r\n\tr = []\r\n\tfor f in file_list:\r\n\t\ts = os.path.splitext(f)\r\n\t\tif s[1] == tail:\r\n\t\t\tr.append(f)\r\n\treturn r\r\n\r\ndef data_augmentation(x, method):\r\n if method == 0:\r\n return np.rot90(x)\r\n if method == 1:\r\n return np.fliplr(x)\r\n if method == 2:\r\n return np.flipud(x)\r\n if method == 3:\r\n return np.rot90(np.rot90(x))\r\n if method == 4:\r\n return np.rot90(np.fliplr(x))\r\n if method == 5:\r\n return np.rot90(np.flipud(x))\r\n\r\n# clear 0.6806, 0.6876, 0.6954\r\n# moire 0.3978, 0.4027, 0.4074\r\ndef calc_meanRGB(img_dirs, tail):\r\n file_list = os.listdir(img_dirs)\r\n file_list = list_filter(file_list, tail)\r\n m = np.zeros((3,))\r\n count = 0\r\n for f in file_list:\r\n count += 1\r\n \r\n img = cv2.imread(img_dirs+f)\r\n img = img.astype(np.float32)/255.0\r\n m += np.mean(img, axis=(0,1))\r\n _m = m/count\r\n print('%d: %s (%f, %f, %f)'%(count, f, _m[0], _m[1], _m[2]), end='\\r')\r\n print(np.round(m/count,4))\r\n\r\ndef crop(x,scale):\r\n shape = x.shape\r\n h = shape[0]\r\n w = shape[1]\r\n h = h-h%scale\r\n w = w-w%scale\r\n return x[0:h,0:w]" } ]