关于mse、sad等测试指标的计算

[jiewjwang]

关于mse、sad等指标的计算，想问下您方便提供一下测试脚本吗，我这边的结果发现对不上，如果能提供，万分感谢

[99991]

Output

   1/1000 - AVG MSE: 3.511 - AVG SAD: 49.633 - 16452523375_08591714cf_o.png
...
  20/1000 - AVG MSE: 4.452 - AVG SAD: 53.613 - 16452523375_08591714cf_o.png
  21/1000 - AVG MSE: 4.352 - AVG SAD: 54.955 - antique-honiton-lace-1182740_1920.png
...
1000/1000 - AVG MSE: 2.192 - AVG SAD: 16.999 - woman-morning-bathrobe-bathroom.png

(Probably a bit better than paper because alpha and composite image are not quantized to uint8 and converted back to float.)

Adobe Composition 1k eval code

import io, math, zipfile, torch, numpy as np
from PIL import Image
from model import AEMatter

def main():
    # Set paths to dataset zip files and model checkpoint here
    adobe_zip_path = "./Adobe_Deep_Matting_Dataset.zip"
    pascal_zip_path = "./PascalVOC2012.zip"
    state_dict = torch.load("./AEMFIX.ckpt", map_location="cpu")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = AEMatter().to(device)
    model.load_state_dict(state_dict["model"])
    eval_model(model, device, adobe_zip_path, pascal_zip_path)

def load_from_zip(path, zf, mode):
    return Image.open(io.BytesIO(zf.read(path))).convert(mode)

def eval_model(model, device, adobe_zip_path, pascal_zip_path):
    model.eval()

    adobe_zip = zipfile.ZipFile(adobe_zip_path)
    pascal_zip = zipfile.ZipFile(pascal_zip_path)

    fg_path = "Combined_Dataset/Test_set/test_fg_names.txt"
    bg_path = "Combined_Dataset/Test_set/test_bg_names.txt"

    fg_names = adobe_zip.read(fg_path).decode("utf-8").splitlines()
    bg_names = adobe_zip.read(bg_path).decode("utf-8").splitlines()

    mses = []
    sads = []
    for i_fg, fg_name in enumerate(fg_names):
        path = f"Combined_Dataset/Test_set/Adobe-licensed images/alpha/{fg_name}"
        alpha_gt = load_from_zip(path, adobe_zip, "L")
        path = f"Combined_Dataset/Test_set/Adobe-licensed images/fg/{fg_name}"
        fg = load_from_zip(path, adobe_zip, "RGB")

        fg = np.float32(fg) / 255.0
        alpha_gt = np.float32(alpha_gt) / 255.0

        for i_trimap in range(20):
            i = i_fg * 20 + i_trimap
            bg_name = bg_names[i]
            trimap_name = fg_name.replace(".png", f"_{i_trimap}.png")

            path = f"Combined_Dataset/Test_set/Adobe-licensed images/trimaps/{trimap_name}"
            trimap = load_from_zip(path, adobe_zip, "L")
            path = f"VOCdevkit/VOC2012/JPEGImages/{bg_name}"
            bg = load_from_zip(path, pascal_zip, "RGB")

            # Resize and crop background to foreground size
            h, w = alpha_gt.shape
            bw, bh = bg.size
            scale = max(w / bw, h / bh)
            if scale > 1:
                bg_w = math.ceil(scale * bw)
                bg_h = math.ceil(scale * bh)
                bg = bg.resize((bg_w, bg_h), Image.BICUBIC)
                bg = bg.crop((0, 0, w, h))

            bg = np.float32(bg) / 255.0
            trimap = np.float32(trimap)

            # Composite foreground and background
            a = alpha_gt[:, :, np.newaxis]
            image = a * fg + (1 - a) * bg

            # Pad image and trimap
            pad_h = -h % 32
            pad_w = -w % 32
            pad_h1 = pad_h // 2
            pad_h2 = pad_h - pad_h1
            pad_w1 = pad_w // 2
            pad_w2 = pad_w - pad_w1

            is_fg = trimap == 255
            is_bg = trimap == 0
            is_unknown = ~(is_fg | is_bg)

            trimap = np.array([is_bg, is_unknown, is_fg])

            image = image.transpose(2, 0, 1)

            padding = ((0, 0), (pad_h1, pad_h2), (pad_w1, pad_w2))

            image = np.pad(image, padding, "reflect")
            trimap = np.pad(trimap, padding, "reflect")

            trimap = torch.from_numpy(trimap).unsqueeze(0).to(device)
            image  = torch.from_numpy(image).unsqueeze(0).to(device)

            # Inference
            with torch.no_grad():
                alpha = model(image, trimap).detach().cpu().numpy()[0, 0]

            # Crop alpha to remove padding
            alpha = alpha[pad_h1 : pad_h1 + h, pad_w1 : pad_w1 + w]

            # Fill in known foreground and background from trimap
            alpha[is_bg] = 0
            alpha[is_fg] = 1

            sad = np.sum(np.abs(alpha - alpha_gt))
            sads.append(sad)

            mse = np.mean(np.square(alpha - alpha_gt)[is_unknown])
            mses.append(mse)

            print(f"{i:4}/1000 - AVG MSE: {np.mean(mses) * 1000:.3f} - AVG SAD: {np.mean(sads) / 1000:.3f} - {fg_name}")

if __name__ == "__main__":
    main()