|
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)): |
|
|
|
image_array = pil_image_as_numpy_array(image) |
|
image = image_array.astype(np.float32) / 255. |
|
|
|
|
|
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): |
|
|
|
|
|
img_file_list = os.listdir(image_folder) |
|
|
|
|
|
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): |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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): |
|
|
|
|
|
predicted_label_indices = np.argmax(y_prediction_encoded, axis=-1) + 1 |
|
|
|
return predicted_label_indices |
|
|
|
def predict_on_train(image): |
|
|
|
sample_image_resized = resize_image(image) |
|
y_pred = ensemble_predict(sample_image_resized) |
|
y_pred = get_predictions(y_pred).squeeze() |
|
|
|
|
|
colors = ['cyan', 'yellow', 'magenta', 'green', 'blue', 'black', 'white'] |
|
|
|
cmap = ListedColormap(colors) |
|
|
|
|
|
fig, ax = plt.subplots() |
|
|
|
|
|
ax.imshow(sample_image_resized) |
|
|
|
|
|
cax = ax.imshow(y_pred, cmap=cmap, vmin=1, vmax=7, alpha=0.5) |
|
|
|
|
|
cbar = plt.colorbar(cax, ticks=np.arange(1, 8)) |
|
cbar.set_ticklabels(['Urban', 'Agriculture', 'Range Land', 'Forest', 'Water', 'Barren', 'Unknown']) |
|
|
|
|
|
image_buffer = io.BytesIO() |
|
plt.savefig(image_buffer, format='png') |
|
image_buffer.seek(0) |
|
image_pil = Image.open(image_buffer) |
|
|
|
|
|
plt.close(fig) |
|
|
|
return image_pil, image_pil |
|
|
|
def predict_on_test(image): |
|
|
|
sample_image_resized = resize_image(image) |
|
y_pred = ensemble_predict(sample_image_resized) |
|
y_pred = get_predictions(y_pred).squeeze() |
|
|
|
|
|
colors = ['cyan', 'yellow', 'magenta', 'green', 'blue', 'black', 'white'] |
|
|
|
cmap = ListedColormap(colors) |
|
|
|
|
|
fig, ax = plt.subplots() |
|
|
|
|
|
ax.imshow(sample_image_resized) |
|
|
|
|
|
cax = ax.imshow(y_pred, cmap=cmap, vmin=1, vmax=7, alpha=0.5) |
|
|
|
|
|
cbar = plt.colorbar(cax, ticks=np.arange(1, 8)) |
|
cbar.set_ticklabels(['Urban', 'Agriculture', 'Range Land', 'Forest', 'Water', 'Barren', 'Unknown']) |
|
|
|
|
|
image_buffer = io.BytesIO() |
|
plt.savefig(image_buffer, format='png') |
|
image_buffer.seek(0) |
|
image_pil = Image.open(image_buffer) |
|
|
|
|
|
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. |
|
''' |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
|
) |
|
|
|
|
|
tab2 = gr.Interface( |
|
fn=predict_on_test, |
|
inputs=gr.Image(), |
|
outputs=gr.Image(), |
|
title='Images with Ground Truth', |
|
description=description, |
|
examples=sample_images |
|
) |
|
|
|
|
|
iface = gr.TabbedInterface([tab1, tab2], |
|
title='Land Cover Segmentation', |
|
tab_names = ['Train','Test']) |
|
|
|
|
|
iface.launch(share=True) |