wbcattplus/evaluation_loop.py at main · apple2373/wbcattplus · 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
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
import torch
import contextlib
from tqdm.auto import tqdm
import numpy as np
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, roc_auc_score
import pandas as pd
from pathlib import Path

# Define a dictionary of evaluation metrics using lambda functions
metrics = {
    "acc": lambda true_labels, predicted_labels: accuracy_score(true_labels, predicted_labels),
    "f1_macro": lambda true_labels, predicted_labels: f1_score(true_labels, predicted_labels, average='macro'),
    "f1_micro": lambda true_labels, predicted_labels: f1_score(true_labels, predicted_labels, average='micro'),
    "f1_weighted": lambda true_labels, predicted_labels: f1_score(true_labels, predicted_labels, average='weighted'),
    "pre_macro": lambda true_labels, predicted_labels: precision_score(true_labels, predicted_labels, average='macro'),
    "pre_micro": lambda true_labels, predicted_labels: precision_score(true_labels, predicted_labels, average='micro'),
    "pre_weighted": lambda true_labels, predicted_labels: precision_score(true_labels, predicted_labels, average='weighted'),
    "rec_macro": lambda true_labels, predicted_labels: recall_score(true_labels, predicted_labels, average='macro'),
    "rec_micro": lambda true_labels, predicted_labels: recall_score(true_labels, predicted_labels, average='micro'),
    "rec_weighted": lambda true_labels, predicted_labels: recall_score(true_labels, predicted_labels, average='weighted'),
}

def evaluate(model, dataloader, criterion=None, scaler=None, true_label_key=None, true_label_key_save=None):
    model.eval()

    attribute_names = dataloader.dataset.attribute_columns
    num_attributes = len(attribute_names)
    attribute_encoders = dataloader.dataset.attribute_encoders
    attribute_decoders = dataloader.dataset.attribute_decoders
    df = dataloader.dataset.df
    image_key = dataloader.dataset.image_key
    if true_label_key is not None and true_label_key_save is None:
        true_label_key_save = f"{true_label_key}_true"
        print('save true *label* (not attribute) as the column name of', f'"{true_label_key_save}"')

    df_true = []  # (N, num_attributes) = (N, 11)
    df_pred_value = []  # (N, num_attributes) = (N, 11)
    df_pred_prob = []  # (N, num_attributes x num_attribute_values) = (N, 31)

    losses = []
    # for some reason, torch.inference_mode() doesn't work with torch.compile()
    # although it says fixed, i still get the error.  https://github.com/pytorch/pytorch/issues/103132
    with torch.no_grad():
        with contextlib.nullcontext() if scaler is None else torch.cuda.amp.autocast():
            # Wrap the dataloader with tqdm to create a progress bar
            for i, data in enumerate(tqdm(dataloader, desc="Evaluate", dynamic_ncols=True)):
                images, image_paths, attributes, attribute_values, indices = data['image'], data['img_path'], data['attributes'], data['attribute_values'], data['idx']
                images, attributes = images.cuda(), attributes.cuda()
                maskb = None
                if "maskb" in data:
                    maskb = data["maskb"].cuda()
                try:
                    model_outputs = model(images)
                except:
                    model_outputs = model(images, maskb)
                # Note that model_outputs is a list of batched attribute softmax logits
                # This means it is a kind of 3-dimensional tensor like model_outputs[att_index][batch_i][att_value_index], where (only) att_value_index has a different length for each attribute.
                # But we can't represent it as a tensor because the length is different from attribute to attribute
                # Also, be aware that the first dimension represents attribute index, not batch index
                if criterion is not None:
                    loss = criterion(model_outputs, attributes.t())
                    losses.append(loss.item())
                # Apply softmax after the loss because the loss function requires logits
                for j in range(num_attributes):
                    # j represents the attribute index
                    model_outputs[j] = torch.softmax(model_outputs[j], dim=1).cpu().detach().numpy()
                # Then loop over each image individually
                # Processing in a batch manner would make the code less readable
                batch_size = len(images)
                for batch_i in range(batch_size):
                    img_path = image_paths[batch_i]
                    true_att_value_indices = attributes[batch_i].cpu().detach().numpy()  # [1, 0, 3, 2, ...]
                    df_row = df.iloc[indices[batch_i].item()]
                    row_base = {image_key: df_row[image_key]}
                    if true_label_key is not None:
                        row_base[true_label_key_save] = df_row[true_label_key]
                    assert Path(row_base[image_key]).name == Path(img_path).name
                    row_att_true = row_base.copy() # e.g., row_att_true["cell_size"] = "small"
                    row_att_pred = row_base.copy()  # e.g., row_att_pred["cell_size"] = "small"
                    row_att_prob = row_base.copy()  # e.g., row_att_prob["cell_size=small"] = 0.9, row_att_prob["cell_size=big"] = 0.1
                    for att_index in range(num_attributes):
                        att_name = attribute_names[att_index]  # e.g., cell_size
                        true_att_value = attribute_values[att_index][batch_i]  # e.g., small
                        att_value_index = true_att_value_indices[att_index]  # e.g., 0
                        assert attribute_encoders[att_name][true_att_value] == att_value_index
                        # Store ground truth value
                        row_att_true[att_name] = true_att_value
                        # Store predicted value
                        p = model_outputs[att_index][batch_i]
                        pred_att_value_index = np.argmax(p)
                        pred_att_value = attribute_decoders[att_name][pred_att_value_index]
                        row_att_pred[att_name] = pred_att_value
                        # Store the binarized version
                        for att_value_idx, p_i in enumerate(p):
                            name = attribute_decoders[att_name][att_value_idx]
                            name = name.replace(' ', '_')
                            row_att_prob[f"{att_name}={name}"] = p_i
                    df_true.append(row_att_true)
                    df_pred_value.append(row_att_pred)
                    df_pred_prob.append(row_att_prob)

    df_true = pd.DataFrame(df_true)
    df_pred_value = pd.DataFrame(df_pred_value)
    df_pred_prob = pd.DataFrame(df_pred_prob)

    per_attribute_metrics = {}
    for att_i, att_name in enumerate(attribute_names):
        for metric_name, metric_func in metrics.items():
            if att_name not in per_attribute_metrics:
                per_attribute_metrics[att_name] = {}
            per_attribute_metrics[att_name][metric_name] = metric_func(df_true[att_name], df_pred_value[att_name])

    overall_metrics = {}
    for metric_name in metrics.keys():
        overall_metrics[metric_name] = np.mean([per_attribute_metrics[att_name][metric_name] for att_name in attribute_names])

    return {
        "overall_metrics": overall_metrics,
        "per_attribute_metrics": per_attribute_metrics,
        "df_true": df_true,
        "df_pred_value": df_pred_value,
        "df_pred_prob": df_pred_prob,
        "loss": np.mean(losses) if len(losses) > 0 else None
    }


def predict(model, dataloader, attribute_encoders, scaler=None):
    model.eval()

    attribute_names = list(attribute_encoders.keys())
    num_attributes = len(attribute_names)
    attribute_decoders = {col: {v: k for k, v in encoding.items()} for col, encoding in attribute_encoders.items()}
    df = dataloader.dataset.df
    image_key = dataloader.dataset.image_key

    df_pred_value = []  # (N, num_attributes) = (N, 11)
    df_pred_prob = []  # (N, num_attributes x num_attribute_values) = (N, 31)

    with torch.no_grad():
        with contextlib.nullcontext() if scaler is None else torch.cuda.amp.autocast():
            for i, data in enumerate(tqdm(dataloader, desc="Predict", dynamic_ncols=True)):
                images, image_paths, indices = data['image'], data['img_path'], data['idx']
                images = images.cuda()

                maskb = data["maskb"].cuda()
                model_outputs = model(images, maskb)

                for j in range(num_attributes):
                    model_outputs[j] = torch.softmax(model_outputs[j], dim=1).cpu().detach().numpy()

                batch_size = len(images)
                for batch_i in range(batch_size):
                    df_row = df.iloc[indices[batch_i].item()]
                    row_att_pred = {image_key: df_row[image_key]}
                    row_att_prob = {image_key: df_row[image_key]}

                    for att_index in range(num_attributes):
                        att_name = attribute_names[att_index]
                        p = model_outputs[att_index][batch_i]
                        pred_att_value_index = np.argmax(p)
                        pred_att_value = attribute_decoders[att_name][pred_att_value_index]
                        row_att_pred[att_name] = pred_att_value

                        for att_value_idx, p_i in enumerate(p):
                            name = attribute_decoders[att_name][att_value_idx]
                            name = name.replace(' ', '_')
                            row_att_prob[f"{att_name}={name}"] = p_i

                    df_pred_value.append(row_att_pred)
                    df_pred_prob.append(row_att_prob)

    df_pred_value = pd.DataFrame(df_pred_value)
    df_pred_prob = pd.DataFrame(df_pred_prob)

    return {
        "df_pred_value": df_pred_value,
        "df_pred_prob": df_pred_prob,
    }