Update app.py

app.py

@@ -74,7 +74,42 @@ def get_predictions(y_prediction_encoded):
 
   return predicted_label_indices
 
-def predict(image):
+def predict_on_train(image):
+    # Steps to get prediction
+    sample_image_resized = resize_image(image)
+    y_pred = ensemble_predict(sample_image_resized)
+    y_pred = get_predictions(y_pred).squeeze()
+
+    # Define your custom colors for each label
+    colors = ['cyan', 'yellow', 'magenta', 'green', 'blue', 'black', 'white']
+    # Create a ListedColormap
+    cmap = ListedColormap(colors)
+
+    # Create a figure
+    fig, ax = plt.subplots()
+
+    # Display the image
+    ax.imshow(sample_image_resized)
+
+    # Display the predictions using the specified colormap
+    cax = ax.imshow(y_pred, cmap=cmap, vmin=1, vmax=7, alpha=0.5)
+
+    # Create colorbar and set ticks and ticklabels
+    cbar = plt.colorbar(cax, ticks=np.arange(1, 8))
+    cbar.set_ticklabels(['Urban', 'Agriculture', 'Range Land', 'Forest', 'Water', 'Barren', 'Unknown'])
+
+    # Convert the figure to a PIL Image
+    image_buffer = io.BytesIO()
+    plt.savefig(image_buffer, format='png')
+    image_buffer.seek(0)
+    image_pil = Image.open(image_buffer)
+
+    # Close the figure to release resources
+    plt.close(fig)
+
+    return image_pil, image_pil
+
+def predict_on_test(image):
     # Steps to get prediction
     sample_image_resized = resize_image(image)
     y_pred = ensemble_predict(sample_image_resized)
@@ -122,7 +157,7 @@ sample_images = get_sample_images('example_images')
 # ).launch(debug=True, share=True)
 
 tab1 = gr.Interface(
-    fn=predict,
+    fn=predict_on_train,
     inputs=gr.Image(label='', type="pil"),
     outputs=[gr.Image(type="pil"), gr.Image(type="pil")],
     title='Images with Ground Truth',
@@ -131,7 +166,7 @@ tab1 = gr.Interface(
 
 # Create the video processing interface
 tab2 = gr.Interface(
-    fn=predict,
+    fn=predict_on_test,
     inputs=gr.File(label=""),
     outputs=gr.File(label=""),
     title='Images with Ground Truth',