bigdl_inference_test.py

# Test dataframe
testDF = NNImageReader.readImages(test_path, sc, resizeH=300, resizeW=300, image_codec=1)
testDF = testDF.withColumn('filename', getFileName('image')).withColumn('label', getLabel('image'))
testPredDF = antbeeModel.transform(testDF).cache()
row = testPredDF.first().asDict()

# showImage function
def showImage(row):
    # Open file
    plt.imshow(Image.open(row['image'][0][5:]))
    # Map prediction to class
    title = 'ants' if row['prediction'] == 1.0 else 'bees'
    plt.title(title)
    
showImage(row)