Update app.py

app.py

@@ -64,7 +64,7 @@ def do_prediction(img):
 
             print(label_p_pred.shape,'label_p_pred')
             print(img_patch.shape, 'img_patch')
-            seg = np.argmax(label_p_pred, axis=2)
+            seg = np.argmax(label_p_pred, axis=3)[0]
 
             print(seg.shape, 'seg')
             seg_color = np.repeat(seg[:, :, np.newaxis], 3, axis=2)