grt/report.py at main · WiseLabCMU/grt · 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
"""Generate Statistical Reports."""

import os
from argparse import ArgumentParser

import yaml
from beartype.typing import Optional
from matplotlib import pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from tqdm import tqdm

import analysis


def _parse(p):
    p.add_argument(
        "-m", "--methods", nargs='+', default=None,
        help="Method comparison sets to generate plots for. Generates all by "
        "default.")
    p.add_argument(
        "-p", "--path", default="results", help="Results directory.")
    p.add_argument(
        "--schema_dir", default="schema", help="Schema base directory.")
    p.add_argument(
        "-s", "--schema", default="base.yaml",
        help="Report schema file.")
    p.add_argument(
        "-o", "--out", default="reports", help="Output directory.")


def _compare(
    results, methods: list[str], method_names: list[str],
    splits: dict[str, str], metric: str = "bev_loss",
    title: Optional[str] = None,
    cmap: str = "coolwarm"
):
    compared = {
        k: results.compare(methods, key=metric, pattern=v)
        for k, v in splits.items()}

    size = 0.75 * len(methods)
    figsize = (1.5 + len(compared) * size, 1.5 + 1.5 * size)
    fig, axs = plt.subplots(3, len(compared), figsize=figsize)
    axs = axs.reshape(3, len(compared))

    analysis.comparison_grid(
        axs, compared, method_names, aspect='auto', cmap=cmap,
        shortnames=all(len(n) <= 10 for n in method_names))

    for k, row in zip(compared, axs.T):
        row[-1].set_xlabel(k)

    if title is not None:
        fig.suptitle(title)
    fig.tight_layout()

    return fig


def _main(args):
    results = analysis.Results(args.path)

    plt.switch_backend('pdf')

    if not args.schema.endswith(".yaml"):
        args.schema = args.schema + ".yaml"
    with open(os.path.join(args.schema_dir, args.schema)) as f:
        cfg = yaml.load(f, Loader=yaml.FullLoader)

    if args.methods is None:
        args.methods = list(cfg["reports"].keys())

    desc = args.schema.replace('.yaml', '') + ":" + ",".join(args.methods)
    for set in tqdm(args.methods, desc=desc):
        spec = cfg["reports"][set]
        with PdfPages(os.path.join(args.out, set + ".pdf")) as document:
            for metric, metric_name in  cfg["metrics"].items():
                fig = _compare(
                    results,
                    methods=spec["methods"],
                    method_names=spec["descriptions"],
                    splits=cfg["splits"],
                    metric=metric,
                    title=f"{spec['name']} / {metric_name}",
                    cmap="coolwarm")
                document.savefig(fig)
                plt.close(fig)


if __name__ == "__main__":
    p = ArgumentParser()
    _parse(p)
    _main(p.parse_args())