File size: 10,110 Bytes
5f0437a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# Copyright (c) 2023 Image Processing Research Group of University Federico II of Naples ('GRIP-UNINA').
#
# All rights reserved.
# This work should only be used for nonprofit purposes.
#
# By downloading and/or using any of these files, you implicitly agree to all the
# terms of the license, as specified in the document LICENSE.txt
# (included in this package) and online at
# http://www.grip.unina.it/download/LICENSE_OPEN.txt

"""
Created in September 2022
@author: fabrizio.guillaro
"""

import logging
import os
import time

import numpy as np
from tqdm import tqdm

import torch
from torch.nn import functional as F

from lib.utils import AverageMeter
from lib.utils import get_confusion_matrix, get_confusion_matrix_1ch
from lib.utils import adjust_learning_rate as default_adjust_learning_rate

def train(epoch, num_epoch, epoch_iters, base_lr, num_iters,
          trainloader, optimizer, model, writer_dict,
          adjust_learning_rate=default_adjust_learning_rate):
    
    # Training
    model.train()
    batch_time = AverageMeter()
    avg_loss = AverageMeter()
    tic = time.time()
    cur_iters = epoch*epoch_iters
    writer = writer_dict['writer']
    global_steps = writer_dict['train_global_steps']
    
    for i_iter, (rgbs, labels) in enumerate(tqdm(trainloader)):
        rgbs = rgbs.cuda()
        labels = labels.long().cuda()

        losses, *_ = model(labels=labels, rgbs=rgbs)
        loss = losses.mean()

        model.zero_grad()
        loss.backward()
        optimizer.step()

        # measure elapsed time
        batch_time.update(time.time() - tic)
        tic = time.time()

        # update average loss
        avg_loss.update(loss.item())

        lr = adjust_learning_rate(optimizer, base_lr, num_iters, i_iter+cur_iters)
      
    print_loss = avg_loss.average()
    msg = 'Epoch: [{}/{}], Time: {:.2f}, ' \
          'lr: {:.6f}, Loss: {:.6f}' .format(
              epoch, num_epoch, batch_time.average(), lr, print_loss)
    logging.info(msg)

    writer.add_scalar('train_loss', print_loss, global_steps)
    writer.add_scalar('learning_rate', lr, global_steps)
    global_steps += 1
    writer_dict['train_global_steps'] = global_steps



def validate(config, testloader, model, writer_dict, valid_set="valid"):

    model.eval()
    avg_loss = AverageMeter()
    confusion_matrix      = np.zeros((config.DATASET.NUM_CLASSES, config.DATASET.NUM_CLASSES))
    confusion_matrix_CONF = np.zeros((config.DATASET.NUM_CLASSES, config.DATASET.NUM_CLASSES))
    

    # PRED metrics
    avg_mse      = AverageMeter()

    avg_mIoU     = AverageMeter()
    avg_p_mIoU   = AverageMeter()
    avg_mIoU_s   = AverageMeter() # smoothed
    avg_p_mIoU_s = AverageMeter() # smoothed
    avg_IoU_1_s  = AverageMeter() # smoothed

    avg_p_F1     = AverageMeter()
    avg_p_F1_s   = AverageMeter()

    # CONF metrics
    c_avg_mse    = AverageMeter()
    c_avg_mIoU   = AverageMeter()
    c_avg_mIoU_s = AverageMeter() # smoothed

    # DET metrics
    avg_det_tpr = AverageMeter()
    avg_det_tnr = AverageMeter()

    
    with torch.no_grad():
        for it, (rgb, label) in enumerate(tqdm(testloader)):
            size = label.size()

            rgb = rgb.cuda()
            label = label.long().cuda()

            losses, pred, conf, det = model(labels=label, rgbs=rgb)
            
            if pred is not None:
                pred = F.upsample(input=pred, size=(size[-2], size[-1]), mode='bilinear')
                pred_prob = F.softmax(pred, dim=1)
                
            if conf is not None:
                conf = F.upsample(input=conf, size=(size[-2], size[-1]), mode='bilinear')
                tcp  = pred_prob[:,1]*(label==1) + pred_prob[:,0]*(label==0)
            
            loss = losses.mean()
            avg_loss.update(loss.item())

            smooth = 1.

            # PRED METRICS
            current_confusion_matrix = get_confusion_matrix(
                label,
                pred,
                size,
                config.DATASET.NUM_CLASSES,
                config.TRAIN.IGNORE_LABEL)

            val_mse = torch.mean((pred_prob[:,1]-label)**2).item()
            avg_mse.update(val_mse)
            confusion_matrix += current_confusion_matrix

            TN = current_confusion_matrix[0, 0]
            FN = current_confusion_matrix[1, 0]
            FP = current_confusion_matrix[0, 1]
            TP = current_confusion_matrix[1, 1]
            pos = current_confusion_matrix.sum(1)    # ground truth label count
            res = current_confusion_matrix.sum(0)    # prediction count
            tp  = np.diag(current_confusion_matrix)  # Intersection part
            
            # mIoU
            IoU_array = (tp / np.maximum(1.0, pos + res - tp))  # Union part
            mean_IoU = IoU_array.mean()                         # mean of the classes IoU            
            avg_mIoU.update(mean_IoU)
            
            # mIoU_s
            IoU_array_smooth = ((tp + smooth)/ (pos + res - tp + smooth))  # Union part
            mean_IoU_smooth = IoU_array_smooth.mean()                      # mean of the classes IoU      
            avg_mIoU_s.update(mean_IoU_smooth)
            avg_IoU_1_s.update(IoU_array_smooth[1])                        # IoU of class 1
            
            # p_mIoU
            p_mIoU = 0.5 * (FN / np.maximum(1.0, FN + TP + TN)) + 0.5 * (FP / np.maximum(1.0, FP + TP + TN))
            avg_p_mIoU.update(np.maximum(mean_IoU, p_mIoU))
            
            # p_mIoU_smooth
            p_mIoU_smooth = 0.5 * ((FN + smooth)/ (FN + TP + TN + smooth)) + 0.5 * ((FP + smooth)/ (FP + TP + TN + smooth))
            avg_p_mIoU_s.update(np.maximum(mean_IoU_smooth, p_mIoU_smooth))

            # p_F1
            F1   = 2 * TP / np.maximum(2 * TP + FN + FP, 1.0)
            p_F1 = 2 * FN / np.maximum(2 * TP + FN + TN, 1.0)
            avg_p_F1.update(np.maximum(F1, p_F1))

            # p_F1_smooth
            F1_s   = (2 * TP + smooth) / (2 * TP + FN + FP + smooth)
            p_F1_s = (2 * FN + smooth) / (2 * TP + FN + TN + smooth)
            avg_p_F1_s.update(np.maximum(F1_s, p_F1_s))


            # CONF metrics
            check_conf = 'conf_head' in config.MODEL.EXTRA.MODULES
            if check_conf:
                current_confusion_matrix_CONF = get_confusion_matrix_1ch(
                    tcp > 0.5,
                    conf,
                    size,
                    config.DATASET.NUM_CLASSES,
                    config.TRAIN.IGNORE_LABEL)
                val_mse = torch.mean((torch.sigmoid(conf[:, 0]) - tcp) ** 2).item()
                c_avg_mse.update(val_mse)
                confusion_matrix_CONF += current_confusion_matrix_CONF

                c_pos = current_confusion_matrix_CONF.sum(1)  # ground truth label count
                c_res = current_confusion_matrix_CONF.sum(0)  # prediction count
                c_tp = np.diag(current_confusion_matrix_CONF)  # Intersection part

                # mIoU (conf)
                c_IoU_array = (c_tp / np.maximum(1.0, c_pos + c_res - c_tp))  # Union part
                c_mean_IoU = c_IoU_array.mean()  # mean of the classes IoU
                c_avg_mIoU.update(c_mean_IoU)

                # mIoU_s (conf)
                c_IoU_array_smooth = ((c_tp + smooth) / (c_pos + c_res - c_tp + smooth))  # Union part
                c_mean_IoU_smooth = c_IoU_array_smooth.mean()  # mean of the classes IoU
                c_avg_mIoU_s.update(c_mean_IoU_smooth)


            # DET metrics
            if det is not None:
                det = det[:,0].cpu().numpy()
            else:
                det = np.max(pred[:,1].cpu().numpy(), axis=(1,2))
            target_det = torch.count_nonzero(label * (label >= 0), (-1, -2)) > 3
            target_det = target_det.cpu().numpy()
            if np.any(target_det>0.5):
                avg_det_tpr.update(np.sum(det[target_det>0.5]>0), np.sum(target_det>0.5))
            if np.any(target_det<0.5):
                avg_det_tnr.update(np.sum(det[target_det<0.5]<0), np.sum(target_det<0.5))


    confusion_matrix = torch.from_numpy(confusion_matrix).cuda()
    confusion_matrix = confusion_matrix.cpu().numpy()
    
    pos = confusion_matrix.sum(1)
    res = confusion_matrix.sum(0)
    tp = np.diag(confusion_matrix)
    pixel_acc = tp.sum() / pos.sum()
    IoU_array = (tp / np.maximum(1.0, pos + res - tp))
    mean_IoU = IoU_array.mean()

    IoU_array_smooth = ((tp + smooth)/ (pos + res - tp + smooth))  # Union part
    mean_IoU_smooth = IoU_array_smooth.mean()                      # mean of the classes IoU      
    
    print_loss = avg_loss.average()
    try:
        bacc = (avg_det_tpr.average()+avg_det_tnr.average())/2
    except:
        bacc = 0
    writer = writer_dict['writer']
    global_steps = writer_dict['valid_global_steps']

    metric_dict = {
        'loss'              : print_loss,
        'mIoU'              : mean_IoU,
        'mIoU_smooth'       : mean_IoU_smooth,
        'avg_mIoU'          : avg_mIoU.average(),
        'avg_mIoU_smooth'   : avg_mIoU_s.average(),
        'avg_det_tpr'       : avg_det_tpr.average(),
        'avg_det_tnr'       : avg_det_tnr.average(),
        'avg_det_bacc'      : bacc,
        'avg_mse'           : avg_mse.average(),
        'avg_IoU_1_smooth'  : avg_IoU_1_s.average(),
        'avg_p-mIoU'        : avg_p_mIoU.average(),
        'avg_p-mIoU_smooth' : avg_p_mIoU_s.average(),
        'avg_p-F1'          : avg_p_F1.average(),
        'avg_p-F1_smooth'   : avg_p_F1_s.average(),
        'pixel_acc'         : pixel_acc,
    }

    if check_conf:
        metric_dict['avg_mse_CONF']         = c_avg_mse.average()
        metric_dict['avg_mIoU_CONF']        = c_avg_mIoU.average()
        metric_dict['avg_mIoU_smooth_CONF'] = c_avg_mIoU_s.average()

    for metric in metric_dict:
        writer.add_scalar(valid_set + '_' + metric, metric_dict[metric], global_steps)

    writer_dict['valid_global_steps'] = global_steps + 1

    return metric_dict, IoU_array, confusion_matrix