Update app.py

app.py

@@ -14,7 +14,6 @@ import os
 import io
 
 REPO_ID = "amosfang/segmentation_u_net"
-image_folder = 'example_images'
 
 def pil_image_as_numpy_array(pilimg):
     img_array = tf.keras.utils.img_to_array(pilimg)
@@ -36,13 +35,14 @@ def load_model_file(filename):
     unet_model = load_model(saved_model_filepath)
     return unet_model
 
-def get_sample_images(image_folder):
+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'.)
-    image_files = [[image_folder +'/' + file] for file in img_file_list if file.lower().endswith(('.jpg', '.jpeg'))]
+    # Filter out only the image files (assuming images have extensions like '.jpg' or '.png')
+    image_files = [[image_folder +'/' + file] for file in img_file_list if file.lower().endswith(format)]
 
     return image_files
 
@@ -75,6 +75,7 @@ def get_predictions(y_prediction_encoded):
   return predicted_label_indices
 
 def predict_on_train(image):
+    
     # Steps to get prediction
     sample_image_resized = resize_image(image)
     y_pred = ensemble_predict(sample_image_resized)
@@ -110,6 +111,7 @@ def predict_on_train(image):
     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)
@@ -146,6 +148,7 @@ def predict_on_test(image):
 
 
 sample_images = get_sample_images('example_images')
+train_images = get_sample_images('train_images')
 
 description= '''
     The DeepGlobe Land Cover Classification Challenge offers the first public dataset containing high resolution 
@@ -167,21 +170,22 @@ description= '''
 #     examples=sample_images
 # ).launch(debug=True, share=True)
 
+# Create the train dataset interface
 tab1 = gr.Interface(
     fn=predict_on_train,
     inputs=gr.Image(),
     outputs=[gr.Image(), gr.Image()],
     title='Images with Ground Truth',
     description=description,
-    examples=sample_images
+    examples=train_images
 )
 
-# Create the video processing interface
+# Create the test dataset interface
 tab2 = gr.Interface(
     fn=predict_on_test,
     inputs=gr.Image(),
     outputs=gr.Image(),
-    title='Images with Ground Truth',
+    title='Images without Ground Truth',
     description=description,
     examples=sample_images
 )