Spaces:
Build error
Build error
Create main.py
Browse files
main.py
ADDED
|
@@ -0,0 +1,39 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import numpy
|
| 2 |
+
import keras
|
| 3 |
+
import gradio
|
| 4 |
+
|
| 5 |
+
# Building the neural network
|
| 6 |
+
model1 = keras.models.Sequential()
|
| 7 |
+
model1.add(keras.layers.InputLayer(input_shape=(101, 636, 1)))
|
| 8 |
+
model1.add(keras.layers.Conv2D(4, (9, 9), activation='relu', padding='same', strides=1))
|
| 9 |
+
model1.add(keras.layers.Conv2D(4, (9, 9), activation='relu', padding='same'))
|
| 10 |
+
model1.add(keras.layers.Conv2D(8, (7, 7), activation='relu', padding='same', strides=1))
|
| 11 |
+
model1.add(keras.layers.Conv2D(8, (7, 7), activation='relu', padding='same'))
|
| 12 |
+
model1.add(keras.layers.Conv2D(16, (5, 5), activation='relu', padding='same'))
|
| 13 |
+
model1.add(keras.layers.Conv2D(16, (5, 5), activation='relu', padding='same', strides=1))
|
| 14 |
+
model1.add(keras.layers.Conv2D(16, (3, 3), activation='relu', padding='same'))
|
| 15 |
+
model1.add(keras.layers.Conv2D(16, (3, 3), activation='relu', padding='same', strides=1))
|
| 16 |
+
model1.add(keras.layers.Conv2D(16, (2, 2), activation='relu', padding='same'))
|
| 17 |
+
model1.add(keras.layers.Conv2D(16, (2, 2), activation='relu', padding='same', strides=1))
|
| 18 |
+
model1.add(keras.layers.UpSampling2D((1, 1)))
|
| 19 |
+
model1.add(keras.layers.Conv2D(16, (2, 2), activation='relu', padding='same'))
|
| 20 |
+
model1.add(keras.layers.UpSampling2D((1, 1)))
|
| 21 |
+
model1.add(keras.layers.Conv2D(8, (3, 3), activation='relu', padding='same'))
|
| 22 |
+
model1.add(keras.layers.UpSampling2D((1, 1)))
|
| 23 |
+
model1.add(keras.layers.Conv2D(4, (7, 7), activation='tanh', padding='same'))
|
| 24 |
+
model1.add(keras.layers.UpSampling2D((1, 1)))
|
| 25 |
+
model1.add(keras.layers.Conv2D(3, (9, 9), activation='tanh', padding='same'))
|
| 26 |
+
|
| 27 |
+
#Loading the weights in the architecture (The file should be stored in the same directory as the code)
|
| 28 |
+
model1.load_weights('modelV13_500trained_1.h5')
|
| 29 |
+
|
| 30 |
+
def predict(mask):
|
| 31 |
+
X = numpy.round((mask/255.0))[numpy.newaxis, :, :, numpy.newaxis]
|
| 32 |
+
v = model1.predict(X)*255
|
| 33 |
+
output = (v - v.min()) / (v.max() - v.min())
|
| 34 |
+
print(output.shape)
|
| 35 |
+
return output[0, :, :, 0], output[0, :, :, 1], output[0, :, :, 2]
|
| 36 |
+
|
| 37 |
+
demo = gradio.Interface(fn=predict, inputs=gradio.Image(image_mode="L", source="canvas", tool="sketch", values=numpy.zeros(636, 101), outputs=[gradio.Image(image_mode="L"), gradio.Image(image_mode="L"), gradio.Image(image_mode="L")])
|
| 38 |
+
|
| 39 |
+
demo.run()
|