Create app.py

app.py

@@ -0,0 +1,94 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from PIL import Image
+from skimage.transform import resize
+import tensorflow as tf
+from tensorflow.keras.models import load_model
+
+import gradio as gr
+import os
+
+REPO_ID = "amosfang/segmentation_u_net"
+
+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():
+    model_dir = snapshot_download(REPO_ID)
+    # saved_model_dir = os.path.join(download_dir, "saved_model")
+    unet_model = load_model(model_dir)
+    return unet_model
+
+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('REPO_ID/train_2024-02-14 11-20-17/base_u_net.0098-acc-0.75-val_acc-0.74-loss-0.79.h5')
+    vgg16_model = load_model('REPO_ID/vgg16_u_net.0092-acc-0.74-val_acc-0.74-loss-0.82.h5')
+    resnet50_model = load_model('REPO_ID/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(image):
+    sample_image_resized = resize_image(image, input_shape)
+    y_pred = ensemble_predict(sample_image_resized)
+    y_pred = get_predictions(y_pred).squeeze()
+
+    # Create a figure without saving it to a file
+    fig, ax = plt.subplots()
+    cax = ax.imshow(y_pred, cmap='viridis', vmin=1, vmax=7)
+
+    # 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
+
+# Specify paths to example images
+sample_images = [['989953_sat.jpg'], ['999380_sat.jpg'], ['988205_sat.jpg']]
+
+# Launch Gradio Interface
+gr.Interface(
+    predict,
+    title='Land Cover Segmentation',
+    inputs=[gr.Image()],
+    outputs=[gr.Image()],
+    examples=sample_images
+).launch(debug=True, share=True)
+
+# Launch the interface
+iface.launch(share=True)