MIDL2020/patch_level.py at master · AIMLab-UBC/MIDL2020 · 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
#!/usr/bin/env python3

from config import get_config, print_usage
from tqdm import tqdm
from PIL import Image
from torch.utils.tensorboard import SummaryWriter
import models.models as models
import numpy as np
import utils.utils as utils
import data.patch_dataset as patch_dataset
import torch
import os


def evaluate(config):
    patches = patch_dataset.PatchDataset(config, apply_color_jitter=False)

    print('{} mode and there are {} patches...'.format(
        config.mode, str(len(patches))))

    data_loader = torch.utils.data.DataLoader(
        patches,
        batch_size=config.batch_size)

    model = models.DeepModel(config, is_eval=True)
    prefix = config.mode + ' Epoch: '

    pred_labels = []
    gt_labels = []
    pred_probs = np.array([]).reshape(0, config.n_subtypes)
    with open(os.path.join(config.dataset_dir, 'results', config.mode.lower() + '_' + config.testing_output_file_name), 'w') as f:
        for data in tqdm(data_loader, desc=prefix, dynamic_ncols=True):
            cur_data, cur_label, patch_infos = data
            with torch.no_grad():
                _, pred_prob = model.forward(cur_data)
                pred_labels += torch.argmax(pred_prob,
                                            dim=1).cpu().numpy().tolist()
                gt_labels += cur_label.cpu().numpy().tolist()
                # write to into distribution output file
                pred_prob = pred_prob.cpu().numpy()
                pred_probs = np.vstack((pred_probs, pred_prob))
                for idx, patch_info in enumerate(patch_infos):
                    f.write('{}\n'.format(patch_info))
                    f.write('{}\n'.format(
                        str(pred_prob[idx]).replace('\n', '')))
                    f.write('---\n')

    utils.compute_metric(gt_labels, pred_labels, pred_probs, verbose=True)


def train(config):
    patches = patch_dataset.PatchDataset(config)

    print('{} mode and there are {} patches...'.format(
        config.mode, str(len(patches))))

    data_loader = torch.utils.data.DataLoader(
        patches,
        batch_size=config.batch_size)

    model = models.DeepModel(config, is_eval=False)

    writer = SummaryWriter(log_dir=os.path.join(config.log_dir, model.name()))
    if not os.path.exists(config.log_dir):
        os.makedirs(config.log_dir)
    if not os.path.exists(config.save_dir):
        os.makedirs(config.save_dir)

    iter_idx = -1
    max_val_acc = float('-inf')
    intv_loss = 0
    for epoch in range(config.epoch):
        prefix = 'Training Epoch {:3d}: '.format(epoch)
        for data in tqdm(data_loader, desc=prefix, dynamic_ncols=True):
            iter_idx += 1

            train_data, train_labels, orig_patch = data
            pred_labels_logits, pred_labels_probs = model.forward(train_data)
            model.optimize_parameters(pred_labels_logits, train_labels)

            intv_loss += model.get_current_errors()

            if iter_idx % config.rep_intv == 0:
                writer.add_scalar('Training CrossEntropyLoss', intv_loss / config.rep_intv,
                                  global_step=iter_idx)
                intv_loss = 0
                val_acc = model.eval(eval_data_ids=patches.eval_data_ids)

                if max_val_acc < val_acc:
                    max_val_acc = val_acc
                    model.save_state(model_id='max_val_acc')

                writer.add_scalar('Validation Accuracy',
                                  val_acc, global_step=iter_idx)


def main(config):
    utils.set_gpus(n_gpus=1)
    if config.mode == 'Training':
        train(config)
    elif config.mode == 'Validation':
        evaluate(config)
    elif config.mode == 'Testing':
        evaluate(config)
    else:
        raise NotImplementedError


if __name__ == '__main__':
    config, unparsed = get_config()

    if len(unparsed) > 0:
        print_usage()
        exit(1)

    main(config)