Update app.py

app.py

@@ -48,20 +48,6 @@ def get_sample_images(image_folder, format=('.jpg', '.jpeg')):
 
     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
@@ -92,8 +78,15 @@ def get_predictions(y_prediction_encoded):
 
 def predict_on_train(image):
 
-    # Steps to get the ground truth image mask
-    ground_truth_mask = get_sample_mask(TRAIN_FOLDER, image)
+    # Extract filename from Gradio input
+    image_filename = gr.get_data()[0].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(TRAIN_FOLDER, mask_filename)
+    ground_truth_mask = Image.open(mask_path)
+    # Resize the mask image
     ground_truth_mask_pil = resize_image(ground_truth_mask)
     
     # Steps to get prediction of the satellite image