Update app.py

app.py

@@ -14,6 +14,9 @@ import os
 import io
 
 REPO_ID = "amosfang/segmentation_u_net"
+TRAIN_FOLDER = 'train_images'
+TEST_FOLDER = 'example_images'
+NUM_CLASSES = 7
 
 def pil_image_as_numpy_array(pilimg):
     img_array = tf.keras.utils.img_to_array(pilimg)
@@ -39,13 +42,26 @@ def get_sample_images(image_folder, format=('.jpg', '.jpeg')):
 
     # Get a list of all files in the folder
     img_file_list = os.listdir(image_folder)
-    img_file_list.sort()
     
-    # Filter out only the image files (assuming images have extensions like '.jpg' or '.png')
+    # Filter out only the image files (assuming images have extensions like '.jpg')
     image_files = [[image_folder +'/' + file] for file in img_file_list if file.lower().endswith(format)]
 
     return image_files
 
+def get_sample_mask(image_folder, image):
+    
+    # Get the filename of the original image
+    image_filename = os.path.basename(image.name)
+    
+    # Construct the filename for the ground truth mask
+    mask_filename = image_filename.replace('_sat.jpg', '_mask.png')
+
+    # Load the ground truth mask
+    mask_path = os.path.join(image_folder, mask_filename)
+    ground_truth_mask = Image.open(mask_path)
+
+    return ground_truth_mask
+
 def ensemble_predict(X_array):
     #
     # Call the predict methods of the unet_model and the vgg16_unet_model
@@ -75,8 +91,12 @@ def get_predictions(y_prediction_encoded):
   return predicted_label_indices
 
 def predict_on_train(image):
+
+    # Steps to get the ground truth image mask
+    ground_truth_mask = get_sample_mask(TRAIN_FOLDER, image)
+    ground_truth_mask_pil = resize_image(ground_truth_mask)
     
-    # Steps to get prediction
+    # Steps to get prediction of the satellite image
     sample_image_resized = resize_image(image)
     y_pred = ensemble_predict(sample_image_resized)
     y_pred = get_predictions(y_pred).squeeze()
@@ -93,10 +113,10 @@ def predict_on_train(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)
+    cax = ax.imshow(y_pred, cmap=cmap, vmin=1, vmax=NUM_CLASSES, alpha=0.5)
 
     # Create colorbar and set ticks and ticklabels
-    cbar = plt.colorbar(cax, ticks=np.arange(1, 8))
+    cbar = plt.colorbar(cax, ticks=np.arange(1, NUM_CLASSES + 1))
     cbar.set_ticklabels(['Urban', 'Agriculture', 'Range Land', 'Forest', 'Water', 'Barren', 'Unknown'])
 
     # Convert the figure to a PIL Image
@@ -108,7 +128,7 @@ def predict_on_train(image):
     # Close the figure to release resources
     plt.close(fig)
 
-    return image_pil, image_pil
+    return ground_truth_mask_pil, image_pil
 
 def predict_on_test(image):
     
@@ -147,8 +167,8 @@ def predict_on_test(image):
     return image_pil
 
 
-sample_images = get_sample_images('example_images')
-train_images = get_sample_images('train_images')
+train_images = get_sample_images(TRAIN_FOLDER)
+sample_images = get_sample_images(TEST_FOLDER)
 
 description= '''
     The DeepGlobe Land Cover Classification Challenge offers the first public dataset containing high resolution