kittendev commited on
Commit
5ccaec8
1 Parent(s): b3c984a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -165,6 +165,8 @@ if __name__ == '__main__':
165
  else:
166
  pred_mask = np.argmax(preds, axis=-1)
167
 
 
 
168
  return pred_mask.astype(np.uint8)
169
 
170
 
 
165
  else:
166
  pred_mask = np.argmax(preds, axis=-1)
167
 
168
+ pred_mask = decode_from_colormap(pred_mask, dataset.colors)[..., ::-1]
169
+
170
  return pred_mask.astype(np.uint8)
171
 
172