Update app.py

app.py

@@ -36,7 +36,7 @@ def inference(image, chosen_model):
     color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) # height, width, 3
     palette = np.array([[0, 0, 0],[255, 255, 255]])
     for label, color in enumerate(palette):
-        color_seg[seg == label, :] = color
+        color_seg[seg == label] = color
     # Convert to BGR
     color_seg = color_seg[..., ::-1]
     img = np.array(image) * 0.5 + color_seg * 0.5