Update main.py

main.py

@@ -1,6 +1,7 @@
 import numpy
 import keras
 import gradio
+import matplotlib.pyplot
 
 # Building the neural network
 model1 = keras.models.Sequential()
@@ -34,14 +35,16 @@ def scale(im, nR, nC):
   return numpy.array([[ im[int(nR0 * r / nR)][int(nC0 * c / nC)]  
              for c in range(nC)] for r in range(nR)])
 
+
 def predict(mask):
-    print(mask)
+    # Get the color map by name:
+    cm = matplotlib.pyplot.get_cmap('RdBu')
+
     X = scale(mask, 101, 636)
     X = numpy.round(X/255.0)[numpy.newaxis, :, :, numpy.newaxis]    
-    v = model1.predict(X)*255
-    output = (v - v.min()) / (v.max() - v.min())
-    print(output.shape)
-    return output[0, :, :, 0], output[0, :, :, 1], output[0, :, :, 2]
+    v = model1.predict(X)
+    output = ((v / max(v.max(), -v.min()))+1.0)/2.0
+    return cm(output[0, :, :, 0]), cm(output[0, :, :, 1]), cm(output[0, :, :, 2])
 
 with gradio.Blocks() as demo:
 
@@ -59,11 +62,10 @@ with gradio.Blocks() as demo:
     mask = gradio.Image(image_mode="L", source="canvas", label="microstructure")    
 
     btn = gradio.Button("Run!", variant="primary")
-    exx = gradio.Image(image_mode="L", label="ε-xx") 
-    eyy = gradio.Image(image_mode="L", label="ε-yy") 
-    exy = gradio.Image(image_mode="L", label="ε-xy")
+    exx = gradio.Image(label="ε-xx") 
+    eyy = gradio.Image(label="ε-yy") 
+    exy = gradio.Image(label="ε-xy")
 
     btn.click(fn=predict, inputs=[mask], outputs=[exx, eyy, exy])
 
-demo.launch()
-
+demo.launch()