Add files via upload

Author: weitingchen83

README.md

@@ -0,0 +1,40 @@
+# PMS-Net: Robust Haze Removal Based on Patch Map for Single Images
+
+Wei-Ting Chen, Jian-Jiun Ding, Sy-Yen Kuo
+
+The IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2019, pp. 11681-11689
+
+In this paper, we proposed a novel haze removal algorithm based on a new feature called the patch map. Conventional patch-based haze removal algorithms (e.g. the Dark Channel prior) usually performs dehazing with a fixed patch size. However, it may produce several problems in recovered results such as oversaturation and color distortion. Therefore, in this paper, we designed an adaptive and automatic patch size selection model called the Patch Map Selection Network (PMS-Net) to select the patch size corresponding to each pixel. This network is designed based on the convolutional neural network (CNN), which can generate the patch map from the image to image. Experimental results on both synthesized and real-world hazy images show that, with the combination of the proposed PMS-Net, the performance in haze removal is much better than that of other state-of-the-art algorithms and we can address the problems caused by the fixed patch size.
+proposed algorithm performs favorably against the stateof-the-art algorithms.
+
+
+# Setup and environment
+
+To generate the recovered result you need:
+1. Python 3 
+2. Keras
+3. CUDA >= 9.0
+4. opencv-python
+
+All the hazy inputs and recovered results are put in the "image" and "result" folder, respectively.
+
+Run the test.py to generate recovered results
+
+```
+$ python test.py
+```
+The pretrained model can be downloaded from: https://drive.google.com/drive/folders/12cKqlfT7ckxwgB-xBdU5TF7Cc5-hssOg?usp=sharing
+
+
+# Citations
+Please cite this paper in your publications if it helps your research:    
+
+@inproceedings{chen2019pms,
+  title={PMS-Net: Robust Haze Removal Based on Patch Map for Single Images},
+  author={Chen, Wei-Ting and Ding, Jian-Jiun and Kuo, Sy-Yen},
+  booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
+  pages={11681--11689},
+  year={2019}
+}
+
+

dehaze_patchMap_dehaze.py

@@ -0,0 +1,121 @@
+import numpy as np
+import cv2
+from find_atmosphericLight import find_atmosphericLight
+from find_darkchannel import find_darkchannel
+from matplotlib import pyplot as plt
+
+def dehaze_patchMap(image, omega, patchMap):
+
+    m, n = image.shape[0], image.shape[1]
+
+    transmissionMap = np.ones((m, n))
+    darkchannelMap = np.ones((m, n))
+
+    patchMap = np.ceil(patchMap)
+    #patchMap(find(patchMap<1))=1
+    #patchMap(find(patchMap>120))=120
+    patchMap[patchMap<1]=1
+    patchMap[patchMap>120]=120
+
+    '''patchMap = guided_filter(rgb2gray(image), patchMap, 15, 0.001)
+    patchMap = ceil(patchMap)
+    patchMap(find(patchMap<1))=1
+    patchMap(find(patchMap>120))=120'''
+
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)/255
+    print(type(gray), gray.shape, type(patchMap), patchMap)
+    patchMap = guidedFilter(gray, patchMap, 15, 0.001).astype(np.uint8)
+    patchMap[patchMap<1]=1
+    patchMap[patchMap>120]=120
+
+    '''[patch_size, ~, patchIdx] = unique(patchMap)
+    patchIdx = reshape(patchIdx, m, n)
+    patch_size_num = size(patch_size)'''
+
+    patch_size = np.unique(patchMap)
+
+    '''for i = 1: patch_size_num(1)
+        i
+        dark_channel = find_darkchannel(image, patch_size(i))
+        atmosphere = find_atmosphericLight(image, dark_channel)
+        atmosphere_est = repmat(reshape(atmosphere, [1, 1, 3]), m, n)
+        est_term = image./atmosphere_est
+        tx_estimation = 1-omega*find_darkchannel(est_term, patch_size(i))
+        tx_estimation = reshape(tx_estimation, m, n)
+        patchIdx = patchMap == patch_size(i)
+        transmissionMap(patchIdx) = tx_estimation(patchIdx)
+        darkchannelMap(patchIdx) = dark_channel(patchIdx)     
+
+
+        tx_refine = guided_filter(rgb2gray(image), transmissionMap, 15, 0.001)
+        tx_refine = reshape(tx_refine, m, n)
+        A_predict = find_atmosphericLight(image, darkchannelMap)
+        A_predict = repmat(reshape(A_predict, [1, 1, 3]), m, n)
+        tx = repmat(max(tx_refine, 0.1), [1, 1, 3])
+        recover_result = ((image - A_predict) ./ tx) + A_predict'''
+    
+    image = image / 255
+    for i in range(len(patch_size)):
+        print(str(i), end= ' ')
+        dark_channel = find_darkchannel(image, patch_size[i])
+
+        atmosphere = find_atmosphericLight(image, dark_channel)
+        
+        #atmosphere_est = repmat(reshape(atmosphere, [1, 1, 3]), m, n)
+        #est_term = image./atmosphere_est
+
+        '''atmosphere_est = np.zeros((m, n, 3))
+        atmosphere_est[...,0] = atmosphere[0]
+        atmosphere_est[...,1] = atmosphere[1]
+        atmosphere_est[...,2] = atmosphere[2]
+        est_term = image / atmosphere_est'''
+        est_term = image / atmosphere
+
+        tx_estimation = 1 - omega * find_darkchannel(est_term, patch_size[i])
+        tx_estimation = np.reshape(tx_estimation, (m, n))
+        '''patchIdx = patchMap == patch_size(i)
+        transmissionMap(patchIdx) = tx_estimation(patchIdx)
+        darkchannelMap(patchIdx) = dark_channel(patchIdx)  '''
+        patchIdx = patchMap==patch_size[i]
+        transmissionMap[patchIdx] = tx_estimation[patchIdx]
+        darkchannelMap[patchIdx] = dark_channel[patchIdx]
+
+        '''tx_refine = guided_filter(rgb2gray(image), transmissionMap, 15, 0.001)
+        tx_refine = reshape(tx_refine, m, n)
+        A_predict = find_atmosphericLight(image, darkchannelMap)
+        A_predict = repmat(reshape(A_predict, [1, 1, 3]), m, n)
+        tx = repmat(max(tx_refine, 0.1), [1, 1, 3])
+        recover_result = ((image - A_predict) ./ tx) + A_predict'''
+        tx_refine = guidedFilter(gray, transmissionMap, 15, 0.001) #guided_filter(rgb2gray(image), transmissionMap, 15, 0.001)
+        #tx_refine = reshape(tx_refine, m, n)
+        A_predict = find_atmosphericLight(image, darkchannelMap)
+        #A_predict = repmat(reshape(A_predict, [1, 1, 3]), m, n)
+        #tx = repmat(max(tx_refine, 0.1), [1, 1, 3])
+        tx = np.reshape(tx_refine,(m,n,1))
+        tx[tx < 0.1] = 0.1
+        tx = np.concatenate((tx,tx,tx), axis=2)
+        recover_result = (image - A_predict) / tx + A_predict
+        #plt.imshow(recover_result); plt.show()
+
+    return recover_result * 255, tx
+    
+def guidedFilter(I, p, r, eps):
+    hei, wid = p.shape
+    N = cv2.boxFilter(np.ones((hei, wid)), -1, (r,r))
+
+    meanI = cv2.boxFilter(I, -1, (r,r)) / N
+    meanP = cv2.boxFilter(p, -1, (r,r)) / N
+    corrI = cv2.boxFilter(I * I, -1, (r,r)) / N
+    corrIp = cv2.boxFilter(I * p, -1, (r,r)) / N
+
+    varI = corrI - meanI * meanI
+    covIp = corrIp - meanI * meanP
+
+    a = covIp / (varI + eps)
+    b = meanP - a * meanI
+
+    meanA = cv2.boxFilter(a, -1, (r,r)) / N
+    meanB = cv2.boxFilter(b, -1, (r,r)) / N
+
+    q = meanA * I + meanB
+    return q

find_atmosphericLight.py

@@ -0,0 +1,18 @@
+import numpy as np
+
+def find_atmosphericLight(image, dark_channel):
+
+    m, n = image.shape[0], image.shape[1]
+
+    search_A_pixel = np.floor(m*n*0.01)
+    image_save = np.reshape(image, (m*n, 3))
+    darkchannel_save = np.reshape(dark_channel, m*n)
+
+    saver = np.zeros((1, 3))
+    idx = np.argsort(-darkchannel_save)
+
+    for pixel_idx in range(int(search_A_pixel)):
+        saver = saver + image_save[ idx[pixel_idx], :]
+
+    A = saver / search_A_pixel
+    return A

find_darkchannel.py

@@ -0,0 +1,11 @@
+import cv2
+
+def find_darkchannel(image, patch_win_size):
+
+    patch_win_size = int(patch_win_size)
+    b,g,r = cv2.split(image)
+    dc = cv2.min(cv2.min(r,g),b)
+    kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(patch_win_size,patch_win_size))
+    dark_channel = cv2.erode(dc,kernel)
+
+    return dark_channel

patchMap_predict.py

@@ -0,0 +1,39 @@
+import numpy as np
+import cv2
+from keras.models import load_model
+import scipy.io as sio
+
+
+
+
+
+base_path_hazyImg = 'image/' 
+base_path_result = 'patchMap/'
+imgname = 'waterfall.tif'
+modelDir = 'PMS-Net.h5'
+
+
+print ("Process image: ", imgname)
+
+
+
+hazy_sample = cv2.imread(base_path_hazyImg + imgname)
+hazy_sample = cv2.resize(hazy_sample,(640,480))
+hazy_input = np.reshape(hazy_sample,(1, 480, 640, 3))
+model = load_model(modelDir)
+patchMap = model.predict(hazy_input, verbose = 1)
+patchMap = np.reshape(patchMap,(-1,1))
+patchMap = np.reshape(patchMap,(480, 640))
+patchMap = np.float64(patchMap)
+
+
+imgname = imgname.replace('.tif','')
+
+
+print('saveDir:',base_path_result + imgname + '.mat')
+
+sio.savemat(base_path_result + imgname + '.mat',{"patchMap":patchMap})
+
+
+
+

recover.py

@@ -0,0 +1,30 @@
+import cv2
+import scipy.io as sio
+import numpy as np
+from dehaze_patchMap_dehaze import dehaze_patchMap
+
+hazy_img_dir = 'image/'
+patchMap_dir = 'patchMap/'
+save_dir = 'result/'
+imgname= 'girls.tif'
+
+
+image = cv2.imread(hazy_img_dir + imgname)
+if image.shape[0] != 480 or image.shape[1] != 640:
+    print('resize image tp 640*480')
+    image = cv2.resize(image,(640,480))
+    
+print(imgname)
+patchMapname = imgname[:-4] + '.mat'
+
+patchMap = sio.loadmat(patchMap_dir + patchMapname)
+patchMap = np.array(patchMap['patchMap'])
+        
+recover_result, tx = dehaze_patchMap(image, 0.95, patchMap)
+
+savename_result = save_dir + 'py_recover_' + imgname
+
+cv2.imwrite(savename_result, recover_result)
+
+
+   

test.py

@@ -0,0 +1,28 @@
+import numpy as np
+import cv2
+from keras.models import load_model
+#import scipy.io as sio
+from dehaze_patchMap_dehaze import dehaze_patchMap
+
+base_path_hazyImg = 'image/' 
+base_path_result = 'patchMap/'
+imgname = 'waterfall.tif'
+save_dir = 'result/'
+modelDir = 'PMS-Net.h5'
+
+print ("Process image: ", imgname)
+
+image = cv2.imread(base_path_hazyImg + imgname)
+if image.shape[0] != 480 or image.shape[1] != 640:
+    print('resize image tp 640*480')
+    image = cv2.resize(image,(640,480))
+hazy_input = np.reshape(image,(1, 480, 640, 3))
+model = load_model(modelDir)
+patchMap = model.predict(hazy_input, verbose = 1)
+patchMap = np.reshape(patchMap,(480, 640))
+        
+recover_result, tx = dehaze_patchMap(image, 0.95, patchMap)
+
+savename_result = save_dir + 'py_recover_' + imgname
+
+cv2.imwrite(savename_result, recover_result)