Merge pull request #14 from snucvpip/jisang

Jisang

.gitignore

@@ -1,6 +1,9 @@
 .DS_Store
-result/*
+result/*/post
+result/*/pre
+result/*/demoire
 .ipynb_checkpoints/*
+sota/epll/NIPSGMM/data
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

data/etc/pre_process/pre_process/pre_process.py

@@ -0,0 +1,90 @@
+import numpy as np
+
+from pre_process.align import *
+import os
+import cv2
+import math
+from skimage.feature import peak_local_max
+from PIL import Image
+
+snapshot_info = []
+
+def GetLimitNxNy(img):
+    pim = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    height, width = pim.shape
+    n_pixel = 30
+
+    coordinates = np.array(peak_local_max(pim, min_distance=3))
+    N = coordinates.shape[0]
+
+    d = np.ceil(np.sqrt(width*height/N))
+
+    return tuple(np.array([width, height]) // int(d * n_pixel))
+
+
+def Snapshot(img, resultdir, level=3):
+    #################### EXPLANATION #####################
+    # This function makes snapshots of each level
+    # The level is equivalent of magnification (the higher level is, the more magnified image is)
+    # The nx, ny determines dividing range
+    # the cropped images in neighborhood have same area for post-processing
+    ######################################################
+
+    ################### USAGE EXAMPLE ####################
+    # Snapshot(5, 2, 2)
+    # This usage will divide image in 5 levels.
+    # The maximum number of dividing is 2 in x, y axis
+    # nx = [1, 3, 5]
+    # ny = [1, 3, 5]
+    # level0 = original input image
+    # level1 = 3 pieces in x-axis X 3 pieces in y-axis = 9 pieces
+    # level2 = 5 pieces in x-axis X 5 pieces in y-axis = 25 pieces
+    ######################################################
+    global snapshot_info
+    nx, ny = GetLimitNxNy(img)
+    list_nx = np.ceil(np.linspace(1, nx, level)).astype(int)
+    list_ny = np.ceil(np.linspace(1, ny, level)).astype(int)
+
+    dir_path = resultdir
+
+    if not os.path.exists(dir_path):
+        os.makedirs(dir_path)
+
+    snapshot_info = []
+    for i in range(level):
+
+        width = int(math.ceil(img.shape[1] * 2 / (list_nx[i] + 1)))
+        height = int(math.ceil(img.shape[0] * 2 / (list_ny[i] + 1)))
+
+        num_image = 1
+        stride = (math.floor(width / 2), math.floor(height / 2))
+        y_start, y_end = (0, height)
+
+        for y in range(list_ny[i]):
+            x_start, x_end = (0, width)
+
+            for x in range(list_nx[i]):
+                cv2.imwrite(dir_path + '/' + str(i) + '-' + str(num_image) + '.png', img[y_start:y_end, x_start:x_end])
+                if list_nx[i] >= 2 and x == list_nx[i] - 2:
+                    x_start = img.shape[1] - 1 - width
+                    x_end = img.shape[1] - 1
+                else:
+                    x_start += stride[0]
+                    x_end += stride[0]
+
+                num_image += 1
+
+            if list_ny[i] >= 2 and y == list_ny[i] - 2:
+                y_start = img.shape[0] - 1 - height
+                y_end = img.shape[0] - 1
+            else:
+                y_start += stride[1]
+                y_end += stride[1]
+
+        snapshot_info.append(((list_nx[i], list_ny[i]), (width, height), stride))
+#         print('level{} Complete, nx: {}, ny: {}, width: {}, height: {}'.format(i, list_nx[i], list_ny[i], width, height))
+
+
+def main(source, target, resultdir):
+    img = cv2.imread(source)
+    Snapshot(img, resultdir)

impl/demoire/epll/denoise.py

@@ -100,6 +100,9 @@ def main(target, matfile, DC, resultdir):
     # cleanI = np.array(Image.open(target).convert('RGB'))/255
     cleanI = denoise(target=target, matpath=matpath, DC=DC, convert_type='L')
 
-    filename = os.path.basename(target)
+    img_type = os.path.basename(target).split('.')[-1]
+    filename = ''.join(os.path.basename(target).split('.')[:-1]) + '_' + ('background' if DC else 'moire') + '.' + img_type
+
     resultpath = os.path.join(resultdir, filename)
+
     save_result(cleanI, resultpath)

impl/demoire/epll/epll.py

@@ -95,9 +95,11 @@ def EPLLhalfQuadraticSplit(noiseI, rambda, patchSize, betas, T, I, LogLFunc, GS,
 
     cleanI = cleanI.reshape(noiseI.shape)
 
-    for i in range(cleanI.shape[0]):
-        for j in range(cleanI.shape[1]):
-            cleanI[i][j] = 1 if cleanI[i][j] > 1 else cleanI[i][j]
-            cleanI[i][j] = 0 if cleanI[i][j] < 0 else cleanI[i][j]
+    cleanI[cleanI > 1] = 1
+    cleanI[cleanI < 0] = 0
+#     for i in range(cleanI.shape[0]):
+#         for j in range(cleanI.shape[1]):
+#             cleanI[i][j] = 1 if cleanI[i][j] > 1 else cleanI[i][j]
+#             cleanI[i][j] = 0 if cleanI[i][j] < 0 else cleanI[i][j]
 
     return cleanI, psnr, cost

impl/demoire/remove_background/remove_background.py

@@ -1,29 +1,57 @@
 import numpy as np
 import matplotlib.pyplot as plt
 import os
-import cv2 as cv
+import cv2
+from skimage.feature import peak_local_max
+from sklearn.cluster import KMeans
 
+def dominantColors(img, n_clusters=3):
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+    #reshaping to a list of pixels
+    img = img.reshape((img.shape[0] * img.shape[1], 3))
+
+    #using k-means to cluster pixels
+    kmeans = KMeans(n_clusters, max_iter=1)
+    kmeans.fit(img)
+
+    #the cluster centers are our dominant colors.
+    return kmeans.cluster_centers_
+
+def smoothing(img):
+    pim = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    height, width = pim.shape
+
+    coordinates = np.array(peak_local_max(pim, min_distance=3))
+    N = coordinates.shape[0]
+
+    d = np.sqrt(width*height/N)
+    z = int(d/2)
+
+    im2 = img.copy().astype(np.uint32)
+    for y, x in coordinates:
+    #     if 100 < pim[y, x] < 150: continue
+        im2[y-z:y+z, x-z:x+z] = (im2[y-z:y+z, x-z:x+z] + img[y, x]) / 2
+
+    return im2.astype(np.uint8)
+
 def seg(img):
-    gray = cv.cvtColor(img,cv.COLOR_BGR2GRAY)
-    G = cv.GaussianBlur(gray,(3,3),0.5)
-    rgb = cv.cvtColor(img,cv.COLOR_BGR2RGB)
-    ret, thresh = cv.threshold(G,0,255, cv.THRESH_BINARY_INV+cv.THRESH_OTSU)
+    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
+    G = cv2.GaussianBlur(gray,(3,3),0.5)
+    rgb = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
+    ret, thresh = cv2.threshold(G,0,255, cv2.THRESH_BINARY_INV+cv2.THRESH_OTSU)
 
     I = np.array(rgb)
     x, y, z = I.shape
-    I[thresh <= 254] = [255,255,255]
-#     for i in range(x):
-#         for j in range(y):
-#             if thresh[i][j] <=254:
-#                 I[i][j] = [255,255,255]
-
-    I2 = cv.GaussianBlur(I, (3,3), 0.5) 
+    I[thresh <= 254] = [255,255,255]                
+    I2 = cv2.GaussianBlur(I, (3,3), 0.5) 
     return I2
 
 def main(target, resultdir):
-    img = cv.imread(target)
-    I = seg(img)   
+    img = cv2.imread(target)
+    img = smoothing(img)   
+    I_seg = seg(img)   
     assert os.path.exists(resultdir), print('result directory not exists')
     filename = os.path.basename(target)
-    plt.imsave(os.path.join(resultdir,filename), I)
+    plt.imsave(os.path.join(resultdir, filename), I_seg)
 

impl/moprem.py

@@ -31,7 +31,7 @@ def __init__(self, imgdir):
         self.post_resultdir = os.path.join(self.resultdir, 'post')
 
         self.eval_datadir = self.post_resultdir
-        self.eval_resultdir = os.path.join(self.resultdir, 'eval')
+        self.eval_resultdir = self.resultdir
 
     def pre(self):
         datadir = self.pre_datadir
@@ -44,15 +44,15 @@ def pre(self):
 
         source = os.path.join(datadir, 'source.png')
         target = os.path.join(datadir, 'target.png')
-        print(resultdir)
         start = time.time()
         pre_process.main(source, target, resultdir)
         end = time.time()
         print('  Pre processing time\t\t\t+{:.2f}s'.format(end-start))
         return self
 
     def demoire(self):
-        def epll(background, moire, pooling=False):
+        '''
+        def _epll(background, moire, pooling=False):
             files = next(os.walk(datadir))[2]
             files.sort()
 
@@ -73,36 +73,36 @@ def epll(background, moire, pooling=False):
             else:
                 start = time.time()
                 for f in files:
-                    if f == '0-1.png':
+                    if f != '2-9.png':
                         continue
                     target = os.path.join(datadir, f)
                     # background seperation
                     epll.main(target, background, True, resultdir)
                     # moire seperation
                     epll.main(target, moire, False, resultdir)
                 end = time.time()
-
-            print('  Epll processing time\t\t\t+{:.2f}s'.format(end-start))
-
-
-        def remove_bg():
+            print('  Epll processing time\t\t\t+{:.2f}s'.format(end-start))'''
+
+        def _remove_background():
             files = next(os.walk(datadir))[2]
             files.sort()
 
             start = time.time()
             for f in files:
+                if f != '0-1.png': continue
                 target = os.path.join(datadir, f)
                 remove_background.main(target, resultdir)
             end = time.time()
-            print('  Remove background processing time\t+{:.2f}s'.format(end-start))
+            print('  Demoire processing time\t\t+{:.2f}s'.format(end-start))
 
         datadir = self.demoire_datadir
         resultdir = self.demoire_resultdir
         assert os.path.exists(datadir), print('Demoire: datadir not exists')
         if not os.path.exists(resultdir):
             os.makedirs(resultdir)
-
-        remove_bg()
+
+#        _epll(background='GSModel_8x8_200_2M_noDC_zeromean.mat', moire='GMM_8x8_200_1500.mat', pooling=False)
+        _remove_background()
         return self
 
     def post(self):
@@ -152,7 +152,7 @@ def eval(self):
             if best_ssim < ssim:
                 best_ssim = ssim
                 best_clean = clean
-        print(best_clean)
+                
         eval.main(source, target, best_clean, resultdir)
         end = time.time()
         print('  Eval time\t\t\t\t+{:.2f}s'.format(end-start))
@@ -178,12 +178,8 @@ def eval(self):
     for imgdir in imgs:
         print('\n--------------------------------------------------')
         print('\'{}\''.format(imgdir))
-
         model = MopReM(imgdir)
-        model.pre()
-        model.demoire()
-        model.post()
-        model.eval()
+        model.pre().demoire().post().eval()
 
     end = time.time()
     print('==================================================')

impl/post_process/eval.py

@@ -178,16 +178,20 @@ def main(source, target, clean, resultdir):
     tar = cv2.cvtColor(tar, cv2.COLOR_BGR2RGB)
 
     plt.close()
-    fig, ax = plt.subplots(ncols=4, nrows=2, figsize=(15, 10))
-    ax[0,0].imshow(src), ax[0,0].set_title('Source')
-    ax[0,1].imshow(tar), ax[0,1].set_title('Target')
-    ax[0,2].imshow(np.abs(src - tar)), ax[0,2].set_title(f'MSE:{round(mse_src, 2)}, PSNR:{round(psnr_src, 2)}, SSIM:{round(ssim_src, 2)}')
-    ax[0,3].imshow(diff_src, cmap='gray', vmin=0, vmax=255), ax[0,3].set_title('Difference')
+    fig, ax = plt.subplots(ncols=3, figsize=(15, 10))
+    print(f'    Source\tMSE: {mse_src:.2f}\tPSNR: {psnr_src:.2f}\tSSIM: {ssim_src:.2f}')
+    ax[0].imshow(src), ax[0].set_title('Source')
+    ax[1].imshow(tar), ax[1].set_title('Target')
+    ax[2].imshow(diff_src, cmap='gray', vmin=0, vmax=255), ax[2].set_title('Difference')
+    plt.tight_layout()
+    plt.savefig(os.path.join(resultdir, 'source_result.pdf'), bbox_inches="tight")
 
-    ax[1,0].imshow(cle), ax[1,0].set_title('Clean')
-    ax[1,1].imshow(tar), ax[1,1].set_title('Target')
-    ax[1,2].imshow(np.abs(cle - tar)), ax[1,2].set_title(f'MSE:{round(mse_cle, 2)}, PSNR:{round(psnr_cle, 2)}, SSIM:{round(ssim_cle, 2)}')
-    ax[1,3].imshow(diff_cle, cmap='gray', vmin=0, vmax=255), ax[1,3].set_title('Difference')
+    plt.close()
+    fig, ax = plt.subplots(ncols=3, figsize=(15, 10))
+    print(f'    Clear\tMSE: {mse_cle:.2f}\tPSNR: {psnr_cle:.2f}\tSSIM: {ssim_cle:.2f}')
+    ax[0].imshow(cle), ax[0].set_title('Clean')
+    ax[1].imshow(tar), ax[1].set_title('Target')
+    ax[2].imshow(diff_cle, cmap='gray', vmin=0, vmax=255), ax[2].set_title('Difference')
     plt.tight_layout()
-    plt.savefig(os.path.join(resultdir, 'result.pdf'), bbox_inches="tight")
+    plt.savefig(os.path.join(resultdir, 'clean_result.pdf'), bbox_inches="tight")
 

impl/post_process/post_process.py

@@ -12,6 +12,7 @@ def main(info, datadir, resultdir):
     result_name = 'clean0.png'
     result_path = os.path.join(resultdir, result_name)
     cv2.imwrite(result_path, img_zero)
+    return
 
     for i in range(1, level):
         img_merge = np.empty(shape=(img_height, img_width, 3))

impl/pre_process/gray_compare.py

@@ -0,0 +1,31 @@
+from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity
+import os
+import cv2
+from matplotlib.pyplot import imshow
+import matplotlib.pyplot as plt
+import numpy as np
+
+datadir = './data'
+src_name = 'background.png'
+tar_name = 'moire.png'
+
+# grayscale
+src = cv2.imread(os.path.join(datadir, src_name), cv2.IMREAD_GRAYSCALE)
+tar = cv2.imread(os.path.join(datadir, tar_name), cv2.IMREAD_GRAYSCALE)
+
+mse = mean_squared_error(src, tar)
+psnr = peak_signal_noise_ratio(src, tar)
+ssim, diff = structural_similarity(src, tar, multichannel=True, full=True)
+
+diff = (diff * 255).astype("uint8")
+
+fig, ax = plt.subplots(ncols=4, figsize=(15, 5))
+ax[0].imshow(src, 'gray')
+ax[0].set_title('Source')
+ax[1].imshow(tar, 'gray')
+ax[1].set_title('Target')
+ax[2].imshow(np.abs(src-tar), 'gray')
+ax[2].set_title(f'MSE:{round(mse,2)}, PSNR:{round(psnr,2)}, SSIM:{round(ssim,2)}')
+ax[3].imshow(diff, cmap='gray')
+ax[3].set_title('Difference')
+plt.savefig(os.path.join(datadir, 'result_gray.png'))