Boltuzamaki commited on
Commit
a7ebeb6
·
1 Parent(s): 6840ca3
.gitignore ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ venv
2
+ .venv
3
+ .env
app.py ADDED
@@ -0,0 +1,156 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import tensorflow as tf
3
+ import numpy as np
4
+ from tensorflow.keras.layers import (
5
+ Conv2D,
6
+ MaxPool2D,
7
+ Dropout,
8
+ Conv2DTranspose,
9
+ concatenate,
10
+ )
11
+ import matplotlib.pyplot as plt
12
+
13
+
14
+ class EncoderBlock(tf.keras.layers.Layer):
15
+ def __init__(self, filters, rate=None, pooling=True, **kwargs):
16
+ super(EncoderBlock, self).__init__(**kwargs)
17
+ self.filters = filters
18
+ self.rate = rate
19
+ self.pooling = pooling
20
+ self.conv1 = Conv2D(
21
+ self.filters,
22
+ kernel_size=3,
23
+ strides=1,
24
+ padding="same",
25
+ activation="relu",
26
+ kernel_initializer="he_normal",
27
+ )
28
+ self.conv2 = Conv2D(
29
+ self.filters,
30
+ kernel_size=3,
31
+ strides=1,
32
+ padding="same",
33
+ activation="relu",
34
+ kernel_initializer="he_normal",
35
+ )
36
+ if self.pooling:
37
+ self.pool = MaxPool2D(pool_size=(2, 2))
38
+ if self.rate is not None:
39
+ self.drop = Dropout(rate)
40
+
41
+ def call(self, inputs):
42
+ x = self.conv1(inputs)
43
+ if self.rate is not None:
44
+ x = self.drop(x)
45
+ x = self.conv2(x)
46
+ if self.pooling:
47
+ y = self.pool(x)
48
+ return y, x
49
+ else:
50
+ return x
51
+
52
+ def get_config(self):
53
+ base_config = super().get_config()
54
+ return {
55
+ **base_config,
56
+ "filters": self.filters,
57
+ "rate": self.rate,
58
+ "pooling": self.pooling,
59
+ }
60
+
61
+
62
+ class DecoderBlock(tf.keras.layers.Layer):
63
+ def __init__(self, filters, rate=None, axis=-1, **kwargs):
64
+ super(DecoderBlock, self).__init__(**kwargs)
65
+ self.filters = filters
66
+ self.rate = rate
67
+ self.axis = axis
68
+ self.convT = Conv2DTranspose(
69
+ self.filters, kernel_size=3, strides=2, padding="same"
70
+ )
71
+ self.conv1 = Conv2D(
72
+ self.filters,
73
+ kernel_size=3,
74
+ activation="relu",
75
+ kernel_initializer="he_normal",
76
+ padding="same",
77
+ )
78
+ if rate is not None:
79
+ self.drop = Dropout(self.rate)
80
+ self.conv2 = Conv2D(
81
+ self.filters,
82
+ kernel_size=3,
83
+ activation="relu",
84
+ kernel_initializer="he_normal",
85
+ padding="same",
86
+ )
87
+
88
+ def call(self, inputs):
89
+ X, short_X = inputs
90
+ ct = self.convT(X)
91
+ c_ = concatenate([ct, short_X], axis=self.axis)
92
+ x = self.conv1(c_)
93
+ if self.rate is not None:
94
+ x = self.drop(x)
95
+ y = self.conv2(x)
96
+ return y
97
+
98
+ def get_config(self):
99
+ base_config = super().get_config()
100
+ return {
101
+ **base_config,
102
+ "filters": self.filters,
103
+ "rate": self.rate,
104
+ "axis": self.axis,
105
+ }
106
+
107
+
108
+ # Load the model with custom layers
109
+ unet = tf.keras.models.load_model(
110
+ "final.h5",
111
+ custom_objects={
112
+ "EncoderBlock": EncoderBlock,
113
+ "DecoderBlock": DecoderBlock,
114
+ },
115
+ )
116
+
117
+
118
+ def show_image(image, cmap=None, title=None):
119
+ plt.imshow(image, cmap=cmap)
120
+ if title is not None:
121
+ plt.title(title)
122
+ plt.axis("off")
123
+
124
+
125
+ def predict(image):
126
+ real_img = tf.image.resize(image, [128, 128])
127
+ real_img = real_img / 255.0
128
+ real_img = np.expand_dims(real_img, axis=0)
129
+ pred_mask = unet.predict(real_img).reshape(128, 128)
130
+ real_img = real_img[0]
131
+
132
+ fig, ax = plt.subplots(1, 2, figsize=(10, 5))
133
+
134
+ ax[0].imshow(real_img)
135
+ ax[0].set_title("Original Image")
136
+ ax[0].axis("off")
137
+
138
+ ax[1].imshow(pred_mask, cmap="gray")
139
+ ax[1].set_title("Predicted Mask")
140
+ ax[1].axis("off")
141
+
142
+ plt.tight_layout()
143
+ plt.show()
144
+ return pred_mask
145
+
146
+
147
+ # Create Gradio interface
148
+ iface = gr.Interface(
149
+ fn=predict,
150
+ inputs=gr.Image(type="numpy"),
151
+ outputs=gr.Image(type="numpy"),
152
+ examples=["./images/water_body_11.jpg", "./images/water_body_1011.jpg"],
153
+ title="Water Body Segmentation",
154
+ )
155
+
156
+ iface.launch()
final.h5 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:61feb05325128544a666a846834378e647951bf67302579f307d76dc095d0368
3
+ size 26090856
images/water_body_1011.jpg ADDED
images/water_body_11.jpg ADDED
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ tensorflow==2.9.0
2
+ gradio
3
+ numpy==1.22