adversarial-mri/display.py at main · saeyslab/adversarial-mri · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
import argparse

import numpy as np

import pandas as pd

import matplotlib.pyplot as plt

import seaborn as sns

import re

from pathlib import Path

pattern = re.compile(r"^tv_(?P<name>[^_]+)_(?P<kind>orig|adv)$")

def repl(match):
    name = match.group("name")
    kind = match.group("kind")
    return (
        f"original-{name}"
        if kind == "orig"
        else f"adversarial-{name}"
    )

def create_ridge_plot(results):
    df_long = (
        results
        .rename(columns=lambda c: c.replace('x_', 'original-').replace('y_', 'adversarial-'))
        .rename(columns=lambda c: re.sub(pattern, repl, c))
        .melt(var_name='source_metric', value_name='value')
    )

    df_long[['source', 'metric']] = df_long['source_metric'].str.split('-', expand=True)
    df_long = df_long.drop(columns='source_metric')

    sns.set_theme(style="white", rc={"axes.facecolor": (0, 0, 0, 0)})

    g = sns.FacetGrid(
        df_long,
        row="metric",
        hue="source",
        aspect=4,
        height=1.4,
        sharex=False
    )

    g.map(
        sns.histplot,
        "value",
        bins="fd",
        stat="count",
        element="step",
        fill=True,
        common_norm=False,
        alpha=0.5
    )

    g.map(plt.axhline, y=0, lw=1, clip_on=False)
    g.axes_dict["mse"].set_xscale("log")

    g.set_titles("")
    g.set(yticks=[], ylabel="")
    g.despine(bottom=True, left=True)

    g.add_legend(title="")

    for metric, ax in g.axes_dict.items():
        ax.text(
            -0.02, 0.5,
            metric.upper(),
            transform=ax.transAxes,
            ha="right",
            va="center",
            fontsize=11,
            fontweight="bold",
            clip_on=False
        )

# parser arguments
parser = argparse.ArgumentParser()
parser.add_argument('-out', type=str, default='./out', help='output directory')
parser.add_argument('-model', type=str, default='unet', choices=['unet', 'varnet'], help='model to use for reconstruction')
parser.add_argument('-organ', type=str, default='knee', choices=['knee', 'brain'])
parser.add_argument('-coil', type=str, default='sc', choices=['sc', 'mc'], help='single-coil (sc) or multi-coil (mc)')
parser.add_argument('-shape', type=str, default='line', choices=['line', 'square'], help='artefact type')

args = parser.parse_args()

# load data
outpath = Path(args.out) / args.model / f"{args.coil}_{args.organ}" / args.shape / "scores.csv"
assert outpath.exists(), f'Path does not exist: {outpath}'

print(f"Loading data from {outpath}")
df = pd.read_csv(outpath)

results = df[['x_psnr', 'y_psnr', 'x_mse', 'y_mse', 'x_ssim', 'y_ssim']]
print(results.describe())

# print best and worst examples
values = df['y_psnr'].to_numpy()
idx = np.argsort(abs(values - np.median(values)))
print(f"Median: {idx[:5]}")

idx = np.argsort(values)
print(f"Lowest PSNR: {idx[:5]}")
print(f"Highest PSNR: {idx[-5:]}")

# create plots
create_ridge_plot(df[['x_psnr', 'y_psnr', 'x_mse', 'y_mse', 'x_ssim', 'y_ssim']])
plt.savefig(f'plots/{args.model}-{args.coil}-{args.organ}-ridge.pdf')
plt.close()

create_ridge_plot(df[['tv_psnr_orig', 'tv_psnr_adv', 'tv_mse_orig', 'tv_mse_adv', 'tv_ssim_orig', 'tv_ssim_adv']])
plt.savefig(f'plots/{args.model}-{args.coil}-{args.organ}-ridge-tv.pdf')
plt.close()