final

.gitignore

@@ -0,0 +1,3 @@
+venv
+.venv
+.env

app.py

@@ -0,0 +1,156 @@
+import gradio as gr
+import tensorflow as tf
+import numpy as np
+from tensorflow.keras.layers import (
+    Conv2D,
+    MaxPool2D,
+    Dropout,
+    Conv2DTranspose,
+    concatenate,
+)
+import matplotlib.pyplot as plt
+
+
+class EncoderBlock(tf.keras.layers.Layer):
+    def __init__(self, filters, rate=None, pooling=True, **kwargs):
+        super(EncoderBlock, self).__init__(**kwargs)
+        self.filters = filters
+        self.rate = rate
+        self.pooling = pooling
+        self.conv1 = Conv2D(
+            self.filters,
+            kernel_size=3,
+            strides=1,
+            padding="same",
+            activation="relu",
+            kernel_initializer="he_normal",
+        )
+        self.conv2 = Conv2D(
+            self.filters,
+            kernel_size=3,
+            strides=1,
+            padding="same",
+            activation="relu",
+            kernel_initializer="he_normal",
+        )
+        if self.pooling:
+            self.pool = MaxPool2D(pool_size=(2, 2))
+        if self.rate is not None:
+            self.drop = Dropout(rate)
+
+    def call(self, inputs):
+        x = self.conv1(inputs)
+        if self.rate is not None:
+            x = self.drop(x)
+        x = self.conv2(x)
+        if self.pooling:
+            y = self.pool(x)
+            return y, x
+        else:
+            return x
+
+    def get_config(self):
+        base_config = super().get_config()
+        return {
+            **base_config,
+            "filters": self.filters,
+            "rate": self.rate,
+            "pooling": self.pooling,
+        }
+
+
+class DecoderBlock(tf.keras.layers.Layer):
+    def __init__(self, filters, rate=None, axis=-1, **kwargs):
+        super(DecoderBlock, self).__init__(**kwargs)
+        self.filters = filters
+        self.rate = rate
+        self.axis = axis
+        self.convT = Conv2DTranspose(
+            self.filters, kernel_size=3, strides=2, padding="same"
+        )
+        self.conv1 = Conv2D(
+            self.filters,
+            kernel_size=3,
+            activation="relu",
+            kernel_initializer="he_normal",
+            padding="same",
+        )
+        if rate is not None:
+            self.drop = Dropout(self.rate)
+        self.conv2 = Conv2D(
+            self.filters,
+            kernel_size=3,
+            activation="relu",
+            kernel_initializer="he_normal",
+            padding="same",
+        )
+
+    def call(self, inputs):
+        X, short_X = inputs
+        ct = self.convT(X)
+        c_ = concatenate([ct, short_X], axis=self.axis)
+        x = self.conv1(c_)
+        if self.rate is not None:
+            x = self.drop(x)
+        y = self.conv2(x)
+        return y
+
+    def get_config(self):
+        base_config = super().get_config()
+        return {
+            **base_config,
+            "filters": self.filters,
+            "rate": self.rate,
+            "axis": self.axis,
+        }
+
+
+# Load the model with custom layers
+unet = tf.keras.models.load_model(
+    "final.h5",
+    custom_objects={
+        "EncoderBlock": EncoderBlock,
+        "DecoderBlock": DecoderBlock,
+    },
+)
+
+
+def show_image(image, cmap=None, title=None):
+    plt.imshow(image, cmap=cmap)
+    if title is not None:
+        plt.title(title)
+    plt.axis("off")
+
+
+def predict(image):
+    real_img = tf.image.resize(image, [128, 128])
+    real_img = real_img / 255.0
+    real_img = np.expand_dims(real_img, axis=0)
+    pred_mask = unet.predict(real_img).reshape(128, 128)
+    real_img = real_img[0]
+
+    fig, ax = plt.subplots(1, 2, figsize=(10, 5))
+
+    ax[0].imshow(real_img)
+    ax[0].set_title("Original Image")
+    ax[0].axis("off")
+
+    ax[1].imshow(pred_mask, cmap="gray")
+    ax[1].set_title("Predicted Mask")
+    ax[1].axis("off")
+
+    plt.tight_layout()
+    plt.show()
+    return pred_mask
+
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="numpy"),
+    outputs=gr.Image(type="numpy"),
+    examples=["./images/water_body_11.jpg", "./images/water_body_1011.jpg"],
+    title="Water Body Segmentation",
+)
+
+iface.launch()

final.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:61feb05325128544a666a846834378e647951bf67302579f307d76dc095d0368
+size 26090856

requirements.txt

@@ -0,0 +1,3 @@
+tensorflow==2.9.0
+gradio
+numpy==1.22