Update app.py

app.py

@@ -1,4 +1,3 @@
-import matplotlib.pyplot as plt
 import numpy as np
 from PIL import Image
 from skimage.transform import resize
@@ -7,6 +6,9 @@ from tensorflow.keras.models import load_model
 
 from huggingface_hub import snapshot_download
 
+import matplotlib.pyplot as plt
+from matplotlib.colors import ListedColormap
+
 import gradio as gr
 import os
 import io
@@ -62,13 +64,24 @@ def get_predictions(y_prediction_encoded):
   return predicted_label_indices
 
 def predict(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 colorbar and set ticks and ticklabels
+    cbar = plt.colorbar(ticks=np.arange(1, 8))
+    cbar.set_ticklabels(['Urban', 'Agriculture', 'Range Land', 'Forest', 'Water', 'Barren', 'Unknown'])
+
+
     # Create a figure without saving it to a file
     fig, ax = plt.subplots()
-    cax = ax.imshow(y_pred, cmap='viridis', vmin=1, vmax=7)
+    cax = ax.imshow(y_pred, cmap=cmap, vmin=1, vmax=7)
 
     # Convert the figure to a PIL Image
     image_buffer = io.BytesIO()