import cv2 import pathlib import os import numpy as np from sample.preprocess import sliceAndMeanShift import matplotlib.pyplot as plt ''' Database file structure example: |-- ImageDatabase │   |-- Images | | |---- | | | |-----Original. | | | |--Sliced | | | | |--128x128 | | | | | |-----128-xPix-yPix | | | | | | |-----128-xPix-yPix-Slice. | | | | | | |--Processed | | | | | | | |-----128-xPix-yPix--Processed. | | | | |--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) # Create dir for sliced images and models slicedImagePath = baseImagePath.joinpath("SL") os.mkdir(slicedImagePath) modelsPath = baseImagePath.joinpath("Models") os.mkdir(modelsPath) # Save Slices sizes = [ "128", "256", "512" ] for size in sizes: print("Processing size:", size) # Run processing processedResult = sliceAndMeanShift(imgArr, int(size)) 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)): # ---128-xPix-yPix x = sliceList[i][0] y = sliceList[i][1] sliceArr = sliceList[i][2] if sliceArr.max() != 0: sliceBaseName = "{}-{}-{}-{}-SL".format(baseImageName, size, x, y) sliceBasePath = slicedSizePath.joinpath(sliceBaseName) os.mkdir(sliceBasePath) # Save Sliced Image sliceFilePath = sliceBasePath.joinpath("{}.{}".format(sliceBaseName, fileExt)) cv2.imwrite(str(sliceFilePath), sliceArr) # Create Processed Slice Path sliceProcPath = sliceBasePath.joinpath("PROC") os.mkdir(sliceProcPath) # Segmented image and label image msSegArr = processedList[i][3] msLabImg = processedList[i][4] msSegDir = sliceProcPath.joinpath("{}-{}-{}".format(sliceBaseName, "MSSEG", "PROC")) msSegFileDir = msSegDir.joinpath("{}.{}".format(msSegDir.name, fileExt)) msLabFileDir = msSegDir.joinpath("{}-{}-{}".format(sliceBaseName, "MSLAB", "PROC")) os.mkdir(msSegDir) # Save processed image and label array cv2.imwrite(str(msSegFileDir), msSegArr) np.save(str(msLabFileDir), msLabImg)