SAInfplus/evaluation_SAInf.py at main · mamazi0131/SAInfplus · 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
import os
from dataloader import get_loader, grid_standardization
from utils import setup_seed
import numpy as np
import json
import torch
import pickle
import time
import argparse
from tqdm import tqdm
from metric import stay_detection_evaluation, stay_selection_evaluation
from SAInf import Detection_Threshold_Estimation

os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

dev_id = 0
os.environ['CUDA_VISIBLE_DEVICES'] = str(dev_id)
torch.cuda.set_device(dev_id)


def evaluation(exp_path, model_name, dataloader, start_time, mode='eval'):
    # train parm
    region = config['region']
    base_path = config['base_path']
    max_candidate_grid = config['max_candidate_grid']
    num_grid = config['num_grid']


    grid_feature_path = os.path.join(base_path, 'grid_feature.pkl')
    grid_feature = pickle.load(open(grid_feature_path, 'rb'))
    grid_feature = grid_standardization(grid_feature)

    grid_feature_map = {gid: feature for gid, feature in enumerate(grid_feature)}

    model_path = os.path.join(exp_path, 'model', model_name)

    model = torch.load(model_path, map_location="cuda:{}".format(dev_id))['model']  # model.to()包含inplace操作，不需要对象承接

    model.eval()

    binary_classification_label_list = []
    whether_stay_pred_list = []
    whether_stay_pred_prob_list = []

    where_stay_label_list = []
    where_stay_pred_list = []

    # unpadding_pair_num = 0
    # stay_pair_num = 0
    # candidate_region_num_list = []
    # detect_stay_pair_num = 0
    # detect_real_stay_pair_num = 0

    for idx, batch in tqdm(enumerate(dataloader)):
        context_data, camera_traj_data, camera_assign_mat, stay_label, \
        candidate_region, camera_pair_feature, candidate_region_feature = batch

        context_data, camera_traj_data, camera_assign_mat, camera_pair_feature = \
            context_data.cuda(), camera_traj_data.cuda(), camera_assign_mat.cuda(), camera_pair_feature.cuda()

        camera_pair_speed, binary_classification_label, \
        stay_query_pair_rep, where_stay_label, pool_mask = model(camera_traj_data,
                                                                 camera_assign_mat,
                                                                 context_data,
                                                                 candidate_region,
                                                                 grid_feature_map,
                                                                 camera_pair_feature,
                                                                 candidate_region_feature,
                                                                 stay_label)

        whether_stay_pred = model.stay_evenet_detection(camera_pair_speed).long()
        binary_classification_label = binary_classification_label.float().cpu()

        whether_stay_pred_prob_list.append(whether_stay_pred.squeeze(-1).cpu().detach().numpy())
        whether_stay_pred_list.append(whether_stay_pred.squeeze(-1).cpu().detach().numpy())
        binary_classification_label_list.append(binary_classification_label)

        where_stay_pred = model.where_stay_head(stay_query_pair_rep)
        where_stay_pred = torch.sigmoid(where_stay_pred)
        where_stay_pred = where_stay_pred * pool_mask

        stay_pair_idx = torch.where(binary_classification_label==1)[0].numpy()

        real_stay_in_where_stay_pred = where_stay_pred[stay_pair_idx].squeeze(-1)  # num_pair*256*1
        real_stay_in_where_stay_label = where_stay_label[stay_pair_idx]  # num_pair*256

        real_stay_in_where_stay_pred = real_stay_in_where_stay_pred.cpu().detach().numpy()
        real_stay_in_where_stay_label = real_stay_in_where_stay_label.float().cpu().detach().numpy()

        where_stay_pred_list.append(real_stay_in_where_stay_pred)
        where_stay_label_list.append(real_stay_in_where_stay_label)

        if mode == 'detail':
            candidate_region_length = []
            for i in range(candidate_region.shape[0]):
                for j in range(candidate_region.shape[1]):
                    if camera_assign_mat[i][j] != num_grid and camera_assign_mat[i][j + 1] != num_grid:
                        candidate_region_set = candidate_region[i][j]
                        candidate_region_length.append(candidate_region_set[candidate_region_set != num_grid].shape[0])
            for i in stay_pair_idx:
                print(where_stay_pred[i][0:candidate_region_length[i]].reshape(-1))
                print(where_stay_label[i][0:candidate_region_length[i]])

    binary_classification_label = np.concatenate(binary_classification_label_list)
    whether_stay_pred = np.concatenate(whether_stay_pred_list)
    whether_stay_pred_prob = np.concatenate(whether_stay_pred_prob_list)

    where_stay_label = np.concatenate(where_stay_label_list)
    where_stay_pred = np.concatenate(where_stay_pred_list)

    acc, auc, precision, recall, f1 = stay_detection_evaluation(binary_classification_label, whether_stay_pred, whether_stay_pred_prob)
    print('acc:{:.4f} auc:{:.4f}'.format(acc, auc))
    print('precision:{:.4f} recall:{:.4f} f1:{:.4f}'.format(precision, recall, f1))

    hit1, hit3, hit5 = stay_selection_evaluation(where_stay_label, where_stay_pred)

    print('hit@1:{:.4f} hit@3:{:.4f} hit@5:{:.4f} SUM:{:.4f}'.format(hit1, hit3, hit5, hit1 + hit3 + hit5))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--region', default='')
    parser.add_argument('--exp_path', default='')
    opt = parser.parse_args()

    assert opt.exp_path.split('/')[-1].split('_')[1] == opt.region, '### region is not equal ###'

    exp_path = opt.exp_path

    config = json.load(open('config/region_{}.json'.format(opt.region), 'r'))

    base_path = config['base_path']
    batch_size = config['batch_size']
    sequence_min_len = config['min_len']
    sequence_max_len = config['max_len']
    num_worker = config['num_worker']
    candidate_threshold = config['candidate_threshold']
    num_grid = config['num_grid']
    max_candidate_grid = config['max_candidate_grid']
    seed = config['random_seed']

    # 设置随机种子
    setup_seed(seed)

    train_loader, val_loader, test_loader = get_loader(base_path, batch_size, num_worker, sequence_min_len,
                                                       sequence_max_len, candidate_threshold, num_grid, max_candidate_grid)


    start_time = time.time()
    # log_path = os.path.join(exp_path, 'evaluation')
    # sys.stdout = Logger(log_path, start_time, stream=sys.stdout)  # record log
    # sys.stderr = Logger(log_path, start_time, stream=sys.stderr)  # record error
    print(exp_path)
    for _,_,files in os.walk(os.path.join(exp_path,'model')):
        for model_name in files:
            print(model_name)
            evaluation(exp_path, model_name, test_loader, start_time, mode='eval')