fcn_predict.py 文件源码

def predict(testImgPath, imgDir, clipSize, net):

    # get file name
    pathArr = testImgPath.split('/')
    tmpFileName = pathArr[len(pathArr) - 1]
    filename = os.path.splitext(tmpFileName)[0]

    # preprocess image
    processedImg = preprocessImg(testImgPath, clipSize)

    # reshape image to be put into data layer
    # shape for input (data blob is N x C x H x W)
    net.blobs['data'].reshape(1, *processedImg.shape)
    print('Predicting...')
    net.blobs['data'].data[...] = processedImg

    # run net and take argmax for prediction
    t = time.process_time()
    net.forward()
    elapsed_time = time.process_time() - t
    out = net.blobs['score'].data[0].argmax(axis=0)
    print("Prediction time: %.3f" % elapsed_time)

    print('Saving...')
    savePrediction(imgDir, out, filename, testImgPath)
    print('Done processing image ' + filename)

    return elapsed_time


# @SUMMARY  : saves output of neural network into four different formats describe above
# @PARAM    : (imgDir) image target directory
# @PARAM    : (out) output of neural network
# @PARAM    : (filename) to save as
# @PARAM    : (testImgPath) for segmentation