FIX: mse bug when unknown area is none
@@ -11,5 +11,9 @@ def sad(alpha, trimap, pred_alpha):
def mse(alpha, trimap, pred_alpha):
assert (pred_alpha[trimap == 0] == 0).all()
assert (pred_alpha[trimap == 255] == 1).all()
- mse = ((pred_alpha - alpha)**2).sum() / (trimap == 128).sum()
+ weight_sum = (trimap == 128).sum()
+ if weight_sum != 0:
+ mse = ((pred_alpha - alpha)**2).sum() / weight_sum
+ else:
+ mse = 0
return mse
