Update app.py

app.py

@@ -31,18 +31,28 @@ shapes = {
     "rectangle": rectangle,
 }
 
+
+
 def randomize(selection):
     index = random.randint(0, 9)
     mask = 255*numpy.array(json.loads(shapes[selection]['Defects'][index]))
     v = numpy.array(json.loads(shapes[selection]['Strain'][index]))
-    print(v)
     # Get the color map by name:
     cm = matplotlib.pyplot.get_cmap('RdBu')
     measure = max(v.max(), -v.min())
-    print(measure)
     output = (v / measure)
-    print(cm((numpy.multiply(output[:, :, 0], mask)+1.0)/2.0))
-    return mask, cm((numpy.multiply(output[:, :, 0], mask/255.0)+1.0)/2.0), cm((numpy.multiply(output[:, :, 1], mask/255.0)+1.0)/2.0), cm((numpy.multiply(output[:, :, 2], mask/255.0)+1.0)/2.0)
+
+    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 mask, cm((numpy.multiply(output[:, :, 0], mask/255.0)+1.0)/2.0), cm((numpy.multiply(output[:, :, 1], mask/255.0)+1.0)/2.0), cm((numpy.multiply(output[:, :, 2], mask/255.0)+1.0)/2.0), legend
 
 with gradio.Blocks() as demo:
     selection = gradio.Dropdown(["circle", "crescent", "ellipse", "peanut", "triangle", "rectangle"], label="Select defect shape")
@@ -56,5 +66,6 @@ with gradio.Blocks() as demo:
             eyy = gradio.Image(label="ε-yy")
         with gradio.Column():
             exy = gradio.Image(label="ε-xy")
+    gradio.HTML(label="", value="")
     selection.change(fn=randomize, inputs=[selection], outputs=[mask, exx, eyy, exy])
 demo.launch(debug=True)