import numpy as np from PIL import Image from skimage.transform import resize import tensorflow as tf 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 REPO_ID = "amosfang/segmentation_u_net" TRAIN_FOLDER = 'train_images' TEST_FOLDER = 'example_images' NUM_CLASSES = 7 def pil_image_as_numpy_array(pilimg): # Convert PIL image to numpy array with Tensorflow utils function img_array = tf.keras.utils.img_to_array(pilimg) return img_array def resize_image(image, input_shape=(224, 224, 3)): # Convert to NumPy array and normalize image_array = pil_image_as_numpy_array(image) image = image_array.astype(np.float32) / 255. # Resize the image to 224x224 image_resized = resize(image, input_shape, anti_aliasing=True) return image_resized def load_model_file(filename): model_dir = snapshot_download(REPO_ID) saved_model_filepath = os.path.join(model_dir, filename) unet_model = load_model(saved_model_filepath) return unet_model def get_sample_images(image_folder, format=[('.jpg', '.jpeg'), ('.png')]): # Initialization image_files = [] mask_files = [] # 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') 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 == []: print(image_files) image_files = [[file] for file in image_files] return image_files return [list(pair) for pair in zip(image_files, mask_files)] def ensemble_predict(X_array): # # Call the predict methods of the unet_model and the vgg16_unet_model # to retrieve their predictions. # # Sum the two predictions together and return their results. # You can also consider multiplying a different weight on # one or both of the models to improve performance X_array = np.expand_dims(X_array, axis=0) unet_model = load_model_file('base_u_net.0098-acc-0.75-val_acc-0.74-loss-0.79.h5') vgg16_model = load_model_file('vgg16_u_net.0092-acc-0.74-val_acc-0.74-loss-0.82.h5') resnet50_model = load_model_file('resnet50_u_net.0095-acc-0.79-val_acc-0.76-loss-0.72.h5') pred_y_unet = unet_model.predict(X_array) pred_y_vgg16 = vgg16_model.predict(X_array) pred_y_resnet50 = resnet50_model.predict(X_array) return (pred_y_unet + pred_y_vgg16 + pred_y_resnet50) / 3 def get_predictions(y_prediction_encoded): # Convert predictions to categorical indices predicted_label_indices = np.argmax(y_prediction_encoded, axis=-1) + 1 return predicted_label_indices def predict_on_train(image, mask): # Get the resized image and mask sample_image_resized = resize_image(image) mask_resized = resize_image(mask) # Create a figure fig, ax = plt.subplots() # Display the image ax.imshow(sample_image_resized) # Display the image cax = ax.imshow(mask_resized, alpha=0.5) # Convert the figure to a PIL Image image_buffer = io.BytesIO() plt.savefig(image_buffer, format='png') image_buffer.seek(0) mask_pil = Image.open(image_buffer) # Close the figure to release resources plt.close(fig) # ---------------------------------------------- # Steps to get prediction of the satellite 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 a figure fig, ax = plt.subplots() # Display the image ax.imshow(sample_image_resized) # Display the predictions using the specified colormap 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, NUM_CLASSES + 1)) cbar.set_ticklabels(['Urban', 'Agriculture', 'Range Land', 'Forest', 'Water', 'Barren', 'Unknown']) # Convert the figure to a PIL Image image_buffer = io.BytesIO() plt.savefig(image_buffer, format='png') image_buffer.seek(0) image_pil = Image.open(image_buffer) # Close the figure to release resources plt.close(fig) return mask_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) 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 a figure fig, ax = plt.subplots() # Display the 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) # Create colorbar and set ticks and ticklabels cbar = plt.colorbar(cax, ticks=np.arange(1, 8)) cbar.set_ticklabels(['Urban', 'Agriculture', 'Range Land', 'Forest', 'Water', 'Barren', 'Unknown']) # Convert the figure to a PIL Image image_buffer = io.BytesIO() plt.savefig(image_buffer, format='png') image_buffer.seek(0) image_pil = Image.open(image_buffer) # Close the figure to release resources plt.close(fig) return image_pil train_images = get_sample_images(TRAIN_FOLDER) test_images = get_sample_images(TEST_FOLDER) description= ''' Computer vision deep learning, powered by GPU advancements, plays a potentially key role in urban planning and climate change research. The U-Net model architecture, that is commonly used in semantic segmentation, can be applied to automated land cover classification and seagrass habitat monitoring. Our project, using the DeepGlobe Land Cover Classification Challenge 2018 dataset, trained four models (basic U-Net, VGG16 U-Net, Resnet50 U-Net, and Efficient Net U-Net) and achieved a validation accuracy of approximately 75\% and a dice score of about 0.6 through an ensemble approach on 803 images with segmentation masks (80/20 split). ''' # Create the train dataset interface tab1 = gr.Interface( fn=predict_on_train, inputs=[gr.Image(), gr.Image()], outputs=[gr.Image(label='Ground Truth'), gr.Image(label='Predicted')], title='Images with Ground Truth', description=description, examples=train_images ) # Create the test dataset interface tab2 = gr.Interface( fn=predict_on_test, inputs=gr.Image(), outputs=gr.Image(label='Predicted'), title='Images without Ground Truth', description=description, examples=test_images ) # Create a Multi Interface with Tabs iface = gr.TabbedInterface([tab1, tab2], title='Land Cover Segmentation', tab_names = ['Train','Test']) # Launch the interface iface.launch(share=True)