database.py 4.65 KB
Newer Older
1
2
3
4
import cv2
import pathlib
import os

josh's avatar
josh committed
5
6
7
8
import numpy as np

from sample.preprocess import sliceAndMeanShift 

josh's avatar
josh committed
9
10
import matplotlib.pyplot as plt

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
'''
Database file structure example:

|-- ImageDatabase
│   |-- Images
|   |   |--<autoInc>-<imgExt>-<name> 
|   |   |   |--<autoInc>-<imgExt>-<name>-Original.<imgExt>
|   |   |   |--Sliced
|   |   |   |   |--128x128
|   |   |   |   |   |--<autoInc>-<imgExt>-<name>-128-xPix-yPix
|   |   |   |   |   |   |--<autoInc>-<imgExt>-<name>-128-xPix-yPix-Slice.<imgExt>  
|   |   |   |   |   |   |--Processed
|   |   |   |   |   |   |   |--<autoInc>-<imgExt>-<name>-128-xPix-yPix-<procType>-Processed.<imgExt>
|   |   |   |   |--256x256
|   |   |   |   |--512x512
'''
def initDB(dbLoc: str, name: str) -> bool:
    '''
    Initialized the image database at a given directory
    File Struture:
    |-- ImageDatabase
    │   |-- Data
    '''
    dbLoc = pathlib.Path(dbLoc)
    dbDir = dbLoc.joinpath(name)
    if dbDir.is_dir():
        return False 
    else:
        imagesDir = dbDir.joinpath("Images")
        os.mkdir(dbDir)
        os.mkdir(imagesDir)
        return True



def importAndProcess(imgPath: str, dbPath: str, newImgName: str):
    '''
    Slices and segments an image and imports it into the database
    '''
    imgPath = pathlib.Path(imgPath)
    dbPath  = pathlib.Path(dbPath)

    # Ensure database exists
    if not dbPath.is_dir():
        print("ERROR: No database exists at {}".format(dbPath))
        return False

    # Internal database directory
    imagesPath = dbPath.joinpath("Images")

    # Ensure database was initialized correctly
    if not imagesPath.is_dir():
        print("ERROR - Database Not Initialized: Use 'initDB.py' to init database.")
        return False

    # PROCESS IMAGE

    # Step 1: Read in image
    imgArr  = cv2.imread(str(imgPath))
    fileExt = imgPath.name.split(".")[-1]

    # Auto increment id
    dbIDs = [int(f.name.split("-")[0]) for f in imagesPath.iterdir() ]
    dbID = 0  

    if len(dbIDs) != 0:
        dbID = max(dbIDs) + 1 

    # Create dir for image
    baseImagePath = imagesPath.joinpath("{}-{}-{}-OR".format(dbID, fileExt, newImgName))
    baseImageName = baseImagePath.name
    os.mkdir(baseImagePath)

    # Step 2: Save Original File
    baseImageFilePath = baseImagePath.joinpath("{}.{}".format(baseImageName, fileExt))
    cv2.imwrite(str(baseImageFilePath), imgArr)

josh's avatar
josh committed
88
    # Create dir for sliced images and models
89
90
    slicedImagePath = baseImagePath.joinpath("SL")
    os.mkdir(slicedImagePath)
josh's avatar
josh committed
91
92
    modelsPath = baseImagePath.joinpath("Models")
    os.mkdir(modelsPath)
93
94
95
96
97
    
    # Save Slices
    sizes = [ "128", "256", "512" ]

    for size in sizes:
josh's avatar
josh committed
98
99
100
101
        print("Processing size:", size)
        # Run processing
        processedResult = sliceAndMeanShift(imgArr, int(size))

102
103
104
105
106
107
108
109
110
111
112
113
114
        slicedSizePath = slicedImagePath.joinpath(size)
        os.mkdir(slicedSizePath)

        # Slices and corresponding processed images
        sliceList     = processedResult["slices"][size]
        processedList = processedResult["meanShift"][size]

        for i in range(len(sliceList)):
            # <autoInc>-<imgExt>-<name>-128-xPix-yPix
            x = sliceList[i][0]
            y = sliceList[i][1]
            sliceArr = sliceList[i][2]

josh's avatar
josh committed
115
116
117
118
            if sliceArr.max() != 0:
                sliceBaseName = "{}-{}-{}-{}-SL".format(baseImageName, size, x, y)
                sliceBasePath = slicedSizePath.joinpath(sliceBaseName)
                os.mkdir(sliceBasePath)
119

josh's avatar
josh committed
120
121
122
                # Save Sliced Image
                sliceFilePath = sliceBasePath.joinpath("{}.{}".format(sliceBaseName, fileExt))
                cv2.imwrite(str(sliceFilePath), sliceArr)
123

josh's avatar
josh committed
124
125
126
                # Create Processed Slice Path
                sliceProcPath = sliceBasePath.joinpath("PROC")
                os.mkdir(sliceProcPath)
127

josh's avatar
josh committed
128
129
130
                # Segmented image and label image
                msSegArr = processedList[i][3]
                msLabImg = processedList[i][4]
131

josh's avatar
josh committed
132
133
134
135
                msSegDir = sliceProcPath.joinpath("{}-{}-{}".format(sliceBaseName, 
                                                                    "MSSEG", 
                                                                    "PROC"))
                msSegFileDir = msSegDir.joinpath("{}.{}".format(msSegDir.name, fileExt))
136

josh's avatar
josh committed
137
138
139
                msLabFileDir = msSegDir.joinpath("{}-{}-{}".format(sliceBaseName, 
                                                                    "MSLAB", 
                                                                    "PROC"))
140

josh's avatar
josh committed
141
                os.mkdir(msSegDir)
142

josh's avatar
josh committed
143
144
145
                # Save processed image and label array
                cv2.imwrite(str(msSegFileDir), msSegArr) 
                np.save(str(msLabFileDir), msLabImg)