RSCDNet/preProcess.py at main · deepanshi-s/RSCDNet · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

import os
import cv2
import sys
from skimage.exposure import match_histograms


def claheapproach(img,cliplimit = 2.0,tilesize = (8,8)):
    clahe=cv2.createCLAHE(clipLimit= cliplimit, tileGridSize=tilesize)
    lab=cv2.cvtColor(img, cv2.COLOR_BGR2LAB)  # convert from BGR to LAB color space
    l,a,b=cv2.split(lab)  # split on 3 different channels
    l2=clahe.apply(l)  # apply CLAHE to the L-channel
    lab=cv2.merge((l2,a,b))  # merge channels
    outputimage=cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)  # convert from LAB to BGR
    return outputimage


def preprocess(path, path2):
    subFolders = os.listdir(path)
    for folder in subFolders:
        os.mkdir(os.path.join(path2), folder)
        newPath = os.path.join(path, folder)
        imageFolders = os.listdir(newPath)
        folderA = os.path.join(newPath, imageFolders[0])
        folderB = os.path.join(newPath, imageFolders[1])
        folderLabel = os.path.join(newPath, imageFolders[2])
        os.mkdir(os.path.join(path2, folder, folderA))
        os.mkdir(os.path.join(path2, folder, folderB))
        os.mkdir(os.path.join(path2, folder, folderLabel))
        imageList = os.listdir(folderA)
        for img in imageList:
            imgA = os.path.join(folderA, img)
            imgB = os.path.join(folderB, img)
            label = os.path.join(folderLabel, img)
            imgA = cv2.imread(imgA)                      #reference for matching histograms
            imgB = cv2.imread(imgB)
            label = cv2.imread(label)
            imgBMatched = match_histograms(imgB, imgA, multichannel=True)

            imgA = claheapproach(imgA)
            imgB = claheapproach(imgBMatched)

            cv2.imwrite(os.path.join(path2, folder, folderA), imgA)
            cv2.imwrite(os.path.join(path2, folder, folderB), imgB)
            cv2.imwrite(os.path.join(path2, folder, folderLabel), label)


if __name__ == "__main__":

    #pass path to the root directory of the dataset
    '''
    The dataset structure:
    train
        - A
            - im1.png
            - im2.png
            ...
        - B
            - im1.png
            - im2.png
            ...
        - label
            - im1.png
            - im2.png
            ...
    val
         - A
            - im1.png
            - im2.png
            ...
        - B
            - im1.png
            - im2.png
            ...
        - label
            - im1.png
            - im2.png
            ...
    test
         - A
            - im1.png
            - im2.png
            ...
        - B
            - im1.png
            - im2.png
            ...
        - label
            - im1.png
            - im2.png
            ...
    '''
    if(len(sys.argv)!=2):
        print("Incorrect arguments, pass the path to the dataset, and the path to save the processed dataset")
        sys.exit

    preprocess(sys.argv[1], sys.argv[2])