fffiloni/Video-Matting-Anything

1"""2Reimplement evaluation.mat provided by Adobe in python3Output of `compute_gradient_loss` is sightly different from the MATLAB version provided by Adobe (less than 0.1%)4Output of `compute_connectivity_error` is smaller than the MATLAB version (~5%, maybe MATLAB has a different algorithm)5So do not report results calculated by these functions in your paper.6Evaluate your inference with the MATLAB file `DIM_evaluation_code/evaluate.m`.7 8by Yaoyi Li9"""10 11import scipy.ndimage12import numpy as np13from skimage.measure import label14import scipy.ndimage.morphology15 16 17def gauss(x, sigma):18    y = np.exp(-x ** 2 / (2 * sigma ** 2)) / (sigma * np.sqrt(2 * np.pi))19    return y20 21 22def dgauss(x, sigma):23    y = -x * gauss(x, sigma) / (sigma ** 2)24    return y25 26 27def gaussgradient(im, sigma):28    epsilon = 1e-229    halfsize = np.ceil(sigma * np.sqrt(-2 * np.log(np.sqrt(2 * np.pi) * sigma * epsilon))).astype(np.int32)30    size = 2 * halfsize + 131    hx = np.zeros((size, size))32    for i in range(0, size):33        for j in range(0, size):34            u = [i - halfsize, j - halfsize]35            hx[i, j] = gauss(u[0], sigma) * dgauss(u[1], sigma)36 37    hx = hx / np.sqrt(np.sum(np.abs(hx) * np.abs(hx)))38    hy = hx.transpose()39 40    gx = scipy.ndimage.convolve(im, hx, mode='nearest')41    gy = scipy.ndimage.convolve(im, hy, mode='nearest')42 43    return gx, gy44 45 46def compute_gradient_loss(pred, target, trimap):47 48    pred = pred / 255.049    target = target / 255.050 51    pred_x, pred_y = gaussgradient(pred, 1.4)52    target_x, target_y = gaussgradient(target, 1.4)53 54    pred_amp = np.sqrt(pred_x ** 2 + pred_y ** 2)55    target_amp = np.sqrt(target_x ** 2 + target_y ** 2)56 57    error_map = (pred_amp - target_amp) ** 258    loss = np.sum(error_map[trimap == 128])59 60    return loss / 1000.61 62 63def getLargestCC(segmentation):64    labels = label(segmentation, connectivity=1)65    largestCC = labels == np.argmax(np.bincount(labels.flat))66    return largestCC67 68 69def compute_connectivity_error(pred, target, trimap, step=0.1):70    pred = pred / 255.071    target = target / 255.072    h, w = pred.shape73 74    thresh_steps = list(np.arange(0, 1 + step, step))75    l_map = np.ones_like(pred, dtype=np.float) * -176    for i in range(1, len(thresh_steps)):77        pred_alpha_thresh = (pred >= thresh_steps[i]).astype(np.int)78        target_alpha_thresh = (target >= thresh_steps[i]).astype(np.int)79 80        omega = getLargestCC(pred_alpha_thresh * target_alpha_thresh).astype(np.int)81        flag = ((l_map == -1) & (omega == 0)).astype(np.int)82        l_map[flag == 1] = thresh_steps[i - 1]83 84    l_map[l_map == -1] = 185 86    pred_d = pred - l_map87    target_d = target - l_map88    pred_phi = 1 - pred_d * (pred_d >= 0.15).astype(np.int)89    target_phi = 1 - target_d * (target_d >= 0.15).astype(np.int)90    loss = np.sum(np.abs(pred_phi - target_phi)[trimap == 128])91 92    return loss / 1000.93 94 95def compute_mse_loss(pred, target, trimap):96    error_map = (pred - target) / 255.097    loss = np.sum((error_map ** 2) * (trimap == 128)) / (np.sum(trimap == 128) + 1e-8)98 99    return loss100 101 102def compute_sad_loss(pred, target, trimap):103    error_map = np.abs((pred - target) / 255.0)104    loss = np.sum(error_map * (trimap == 128))105 106    return loss / 1000, np.sum(trimap == 128) / 1000107 108def compute_mad_loss(pred, target, trimap):109    error_map = np.abs((pred - target) / 255.0)110    loss = np.sum(error_map * (trimap == 128)) / (np.sum(trimap == 128) + 1e-8)111 112    return loss113