Update app.py

app.py

@@ -49,12 +49,12 @@ def get_sample_images(image_folder, format=[('.jpg', '.jpeg'), ('.png')]):
     img_file_list = os.listdir(image_folder)
     img_file_list.sort()
     
-    # Filter out only the image files (assuming images have extensions like '.jpg')
+    # 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[0])]
     mask_files = [image_folder +'/' + file for file in img_file_list if file.lower().endswith(format[1])]
 
     if mask_files == []:
-        image_files = [list(file) for file in image_files]
+        image_files = [[file] for file in image_files]
         return image_files
 
     return [list(pair) for pair in zip(image_files, mask_files)]