Update app.py

app.py

@@ -17,6 +17,10 @@ def infer(original_image):
     output_image = model.predict(image)
     output_image = tf.cast((output_image[0, :, :, :] * 255), dtype=np.uint8)
     output_image = Image.fromarray(output_image.numpy())
+    output_image = output_image.reshape(
+        (np.shape(output_image)[0], np.shape(output_image)[1], 3)
+    )
+    output_image = np.uint32(output_image)
     return output_image
 
 iface = gr.Interface(