Update main.py

Better plotting

main.py

@@ -2,6 +2,7 @@ import numpy
 import keras
 import gradio
 import matplotlib.pyplot
+import matplotlib.colors
 
 # Building the neural network
 model1 = keras.models.Sequential()
@@ -28,6 +29,9 @@ model1.add(keras.layers.Conv2D(3, (9, 9), activation='tanh', padding='same'))
 #Loading the weights in the architecture (The file should be stored in the same directory as the code)
 model1.load_weights('modelV13_500trained_1.h5')
 
+# Get the color map by name:
+cm = matplotlib.pyplot.get_cmap('RdBu')
+
 #simple image scaling to (nR x nC) size
 def scale(im, nR, nC):
   nR0 = len(im)     # source number of rows 
@@ -37,15 +41,22 @@ def scale(im, nR, nC):
 
 
 def predict(mask):
-    # Get the color map by name:
-    cm = matplotlib.pyplot.get_cmap('RdBu')
-
-    scaled_mask = numpy.round(scale(mask, 101, 636)/255.0)
+    scaled_mask = numpy.ones((101, 636)) if mask is None else numpy.round(scale(mask, 101, 636)/255.0)
     print(scaled_mask)
     X = scaled_mask[numpy.newaxis, :, :, numpy.newaxis]    
     v = model1.predict(X)
-    output = (v / max(v.max(), -v.min()))
-    return cm((numpy.multiply(output[0, :, :, 0], scaled_mask)+1.0)/2.0), cm((numpy.multiply(output[0, :, :, 1], scaled_mask)+1.0)/2.0), cm((numpy.multiply(output[0, :, :, 2], scaled_mask)+1.0)/2.0)
+    measure = max(v.max(), -v.min())
+    output = (v / measure)
+    legend = "<h2>Strain</h2><table style=\"width:100%\"><tr>"
+    for i in range(11):
+        color = cm(i/10.0)[:3]
+        value = -measure + i*2*measure/10
+        print(sum(list(color)))
+        hex = matplotlib.colors.to_hex(list(color))
+        text_color = "black" if sum(list(color)) > 2.0 else "white"
+        legend = legend + f"<td style=\"background-color: {hex}; color: {text_color}\">{value:+.2e}</td>"
+    legend = legend + "</tr></table>"
+    return cm((numpy.multiply(output[0, :, :, 0], scaled_mask)+1.0)/2.0), cm((numpy.multiply(output[0, :, :, 1], scaled_mask)+1.0)/2.0), cm((numpy.multiply(output[0, :, :, 2], scaled_mask)+1.0)/2.0), legend
 
 with gradio.Blocks() as demo:
 
@@ -66,7 +77,8 @@ with gradio.Blocks() as demo:
     exx = gradio.Image(label="ε-xx") 
     eyy = gradio.Image(label="ε-yy") 
     exy = gradio.Image(label="ε-xy")
+    legend = gradio.HTML(label="", value="")
 
-    btn.click(fn=predict, inputs=[mask], outputs=[exx, eyy, exy])
+    btn.click(fn=predict, inputs=[mask], outputs=[exx, eyy, exy, legend])
 
 demo.launch(debug=True)