surface_prediction.py

# pick random test data sample from one batch
x = random.randint(0, 32 - 1) # default batch size is 32

for i in test_df.as_numpy_iterator():
    img, label = i    
    plt.axis('off')   # remove axes
    plt.imshow(img[x])    # shape from (32, 256, 256, 3) --> (256, 256, 3)
    output = model.predict(np.expand_dims(img[x],0))[0][0]    # getting output; input shape (256, 256, 3) --> (1, 256, 256, 3)
    pred = (output > 0.5).astype('int')
    print("Predicted: ", class_names[pred], '(', output, '-->', pred, ')')    # Picking the label from class_names base don the model output
    print("True: ", class_names[label[x][0].astype('int')])
    break