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" image_folder = 'example_images' def pil_image_as_numpy_array(pilimg): 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): # Get a list of all files in the folder img_file_list = os.listdir(image_folder) # 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'))] return image_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): # 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, 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 sample_images = get_sample_images('example_images') description= ''' The DeepGlobe Land Cover Classification Challenge offers the first public dataset containing high resolution satellite imagery focusing on rural areas. As there are multiple land cover types and high density of annotations, this dataset is more challenging than its counterparts launched in 2018. All satellite images contain RGB pixels, with a pixel resolution of 50 cm. The total size of the total area of the dataset is equivalent to 10716.9 square kilometers. We trained on 803 images and their segmentation masks (with split of 80/20%). For this multilabel segmentation task, we trained 4 models, the basic 4-blocks U-net CNN, VGG16 U-Net, Resnet50 U-net and Efficient Net U-net. Then, I built an ensemble model that achieved a validation accuracy of about 75% and dice score of about 0.6. ''' # Launch Gradio Interface (Single Tab interface) # gr.Interface( # predict, # title='Land Cover Segmentation', # inputs=[gr.Image()], # outputs=[gr.Image()], # examples=sample_images # ).launch(debug=True, share=True) 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 ) # Create the video processing interface tab2 = gr.Interface( fn=predict_on_test, inputs=gr.Image(), outputs=gr.Image(), title='Images with Ground Truth', description=description, examples=sample_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)