data.py 文件源码

def loadImages(datadir, maxDirectoryCount=10, split=0.9):
    for dirPath, dirNames, fileNames in os.walk(datadir):
        fileNames = [f for f in fileNames if not f[0] == '.']
        dirNames[:] = [d for d in dirNames if not d[0] == '.']
        if (maxDirectoryCount != 0):
            fullSizeFileNames = [fileName for fileName in fileNames if fileName.endswith("@2x.png") and (fileName.replace("@2x","") in fileNames)]
            for fullSizeFileName in fullSizeFileNames:
                inputImage = io.imread(dirPath + "/" + fullSizeFileName)
                targetImage = io.imread(dirPath + "/" + fullSizeFileName.replace("@2x",""))
                # print(dirPath + "/" + fullSizeFileName)
                inputSlices, targetSlices = sliceImages(inputImage, targetImage)
                # print("got", len(inputSlices), "input splices and",len(targetSlices),"targetSlices")
                inputImages.extend(inputSlices)
                targetImages.extend(targetSlices)
            maxDirectoryCount -= 1
    x, y = np.asarray(inputImages), np.asarray(targetImages)
    x_train = x[:int(len(x) * split)]
    y_train = y[:int(len(y) * split)]
    x_test = x[int(len(x) * split):]
    y_test = y[int(len(y) * split):]
    # Shuffle training data so that repeats aren't in the same batch
    # x_train, y_train = shuffle(x_train, y_train, random_state=0)
    return (x_train, y_train, x_test, y_test)