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)