StrokeFormer/test_rpn.py at main · angelonazzaro/StrokeFormer · 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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
import csv
import os
from argparse import ArgumentParser

import torch
from lightning import seed_everything
from torchvision.transforms.v2.functional import to_pil_image
from torchvision.utils import make_grid, draw_bounding_boxes
from tqdm import tqdm

from dataset import RegionProposalDataModule
from models import RPN
from utils import build_metrics, compute_metrics, get_device, to_3channel


def test(args):
    seed_everything(args.seed)

    model = RPN.load_from_checkpoint(args.ckpt_path)
    model = model.to(device=get_device())
    model.eval()
    model_name = args.model_name or "_".join(args.ckpt_path.split(os.path.sep)[-1].split("-")[:2])

    datamodule = RegionProposalDataModule(
        paths={
            "test": {
                "scans": args.scans,
                "masks": args.masks,
            }
        },
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        resize_to=args.resize_to
    )

    datamodule.setup(stage="test")
    dataloader = datamodule.test_dataloader()

    metrics = build_metrics(num_classes=args.num_classes, task="region_proposal")

    os.makedirs(args.scores_dir, exist_ok=True)
    model_prediction_dir = os.path.join(args.scores_dir, model_name, "predictions")

    os.makedirs(model_prediction_dir, exist_ok=True)
    global_scores = {}


    for batch_idx, batch in enumerate(tqdm(dataloader, desc="Proposing regions")):
        slices, targets = batch["slices"], batch["targets"]
        targets = [{k: v.to(model.device) for k, v in t.items()} for t in targets]
        slices = slices.to(device=model.device)

        with torch.no_grad():
            proposals = model(slices, None)

        preds = [
            {
                "boxes": p["boxes"].cpu(),
                "scores": p["scores"].cpu(),
                "labels": p["labels"].cpu()
            }
            for p in proposals
        ]
        targets = [
            {
                "boxes": t["boxes"].cpu(),
                "labels": t["labels"].cpu()
            }
            for t in targets
        ]

        # focus the metrics computation on brain regions only
        scores = compute_metrics(preds, targets, metrics, task="region_proposal")

        if len(global_scores) == 0:
            for metric_name in scores.keys():
                if metric_name == "classes":
                    continue
                global_scores[metric_name] = {
                    "ca": scores[metric_name],
                    "n": 1,
                }
        else:
            # CA update rule: (x_n+1 + n * CA_n) / (n + 1)
            for metric_name in global_scores.keys():
                if metric_name == "classes":
                    continue
                curr_ca = global_scores[metric_name]["ca"]
                curr_n = global_scores[metric_name]["n"]
                global_scores[metric_name] = {
                    "ca": (scores[metric_name] + curr_n * curr_ca) / (curr_n + 1),
                    "n": curr_n + 1
                }

        if args.n_predictions > batch_idx:
            if slices.min() != 0 or slices.max() != 1:
                slices = slices - slices.min()
                slices = slices / slices.max()
            ground_truths = [draw_bounding_boxes(to_3channel(img[:1]), target["boxes"], colors="cyan").cpu() for img, target in zip(slices, targets)]
            predictions = [draw_bounding_boxes(to_3channel(img[:1]), proposal["boxes"], colors="red").cpu() for img, proposal in zip(slices, proposals)]
            ground_truths = torch.stack(ground_truths)
            predictions = torch.stack(predictions)

            all_imgs = torch.cat([ground_truths, predictions], dim=0)

            n_pairs = slices.shape[0]
            # interleave reference and reconstruction: ref0, rec0, ref1, rec1, ...
            paired_images = torch.empty((n_pairs * 2, 3, *slices.shape[2:]), dtype=slices.dtype)
            paired_images[0::2] = all_imgs[:n_pairs]
            paired_images[1::2] = all_imgs[n_pairs:]
            grid = make_grid(paired_images, padding=2, pad_value=1.0)
            grid = to_pil_image(grid)
            grid.save(os.path.join(model_prediction_dir, f"{batch_idx}_preds.png"))

    scores_path = os.path.join(args.scores_dir, args.scores_file)
    file_exists = os.path.exists(scores_path)

    global_scores = {m: v["ca"].item() for m, v in global_scores.items()}
    global_scores = {"model": args.model_name, **global_scores}

    with open(scores_path, "a", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=global_scores.keys())
        if not file_exists:
            writer.writeheader()
        writer.writerow(global_scores)


if __name__ == "__main__":
    parser = ArgumentParser()

    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--scans", type=str, required=True)
    parser.add_argument("--masks", type=str, required=True)
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument("--num_workers", type=int, default=0)
    parser.add_argument("--resize_to", nargs=2, type=int, default=None)

    parser.add_argument("--ckpt_path", type=str, required=True)
    parser.add_argument("--model_name", type=str, default=None)
    parser.add_argument("--num_classes", type=int, default=2)

    parser.add_argument("--n_predictions", help="Number of predictions to save", type=int, default=30)
    parser.add_argument("--scores_dir", type=str, default="./rpn_scores")
    parser.add_argument("--scores_file", type=str, default="rpn_scores.csv")

    test(parser.parse_args())