Update app.py

app.py

@@ -223,9 +223,11 @@ def do_prediction(model_name, img):
 
                 img_border = np.zeros((prediction_true.shape[0],prediction_true.shape[1]))
                 img_border[y:y+h, x:x+w] = 1
+                img_border = np.repeat(img_border[:, :, np.newaxis], 3, axis=2)
             else:
                 img_border = np.zeros((prediction_true.shape[0],prediction_true.shape[1]))
                 img_border[:, :] = 1
+                img_border = np.repeat(img_border[:, :, np.newaxis], 3, axis=2)
                 
             return "No numerical output", visualize_model_output(img_border,img, model_name)