Update main.py

main.py

@@ -40,11 +40,12 @@ def predict(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]    
+    scaled_mask = 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()))+1.0)/2.0
-    return cm(output[0, :, :, 0]), cm(output[0, :, :, 1]), cm(output[0, :, :, 2])
+    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)
 
 with gradio.Blocks() as demo:
 
@@ -68,4 +69,4 @@ with gradio.Blocks() as demo:
 
     btn.click(fn=predict, inputs=[mask], outputs=[exx, eyy, exy])
 
-demo.launch()
+demo.launch(debug=True)