forked from sair-lab/AirCode
-
Notifications
You must be signed in to change notification settings - Fork 2
/
validate_gcn.py
168 lines (139 loc) · 5.17 KB
/
validate_gcn.py
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
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import datetime
import logging
import os
import time
import argparse
import yaml
import torch
import torch.distributed as dist
from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
from datasets.utils.build_data import coco_loader
from model.build_model import build_maskrcnn, build_gcn
from datasets.utils.preprocess import warp_batch_data, match_points_clusters
from validate import maskrcnn_inference
from model.graph_models.descriptor_loss import DescriptorLoss
def calculate_f1(maskrcnn_model, gcn_model, loader, configs):
with torch.no_grad():
maskrcnn_model.eval()
gcn_model.eval()
## data cofig
data_config = configs['data']
## superpoint model config
superpoint_model_config = configs['model']['superpoint']
detection_threshold = superpoint_model_config['eval']['detection_threshold']
precisions, recalls, weights = [], [], []
for iter, batch in enumerate(loader):
optimizer.zero_grad()
original_images = batch['image']
original_sizes = [list(img.shape[-2:]) for img in original_images]
_, points_output, maskrcnn_targets, _ = maskrcnn_inference(
maskrcnn_model, batch, use_gpu, 1, data_config, detection_threshold)
warped_batch = warp_batch_data(batch, data_config)
_, warped_points_output, warped_maskrcnn_targets, _ = maskrcnn_inference(
maskrcnn_model, warped_batch, use_gpu, 1, data_config, detection_threshold)
batch_points, batch_descs, connections = match_points_clusters(points_output, maskrcnn_targets['masks'],
warped_points_output, warped_maskrcnn_targets['masks'])
if len(connections) < 2:
print("no object")
continue
batch_points = [points.cuda() for points in batch_points]
batch_descs = [descs.cuda() for descs in batch_descs]
batch_object_descs = gcn_model(batch_points, batch_descs)
connections = torch.stack(connections).cuda()
distances = torch.einsum('nd,dm->nm', descs, descs.t()) # N * N
good_matchs = (distances > dist_thr).float()
num_correct_matches = torch.sum(good_matchs * connections)
num_connections = torch.sum(connections)
recall = num_correct_matches / num_connections
precision = num_correct_matches / torch.sum(good_matchs)
recalls.append(recall)
precisions.append(precision)
weights.append(num_connections)
if(len(weights) == 0):
return 0., 0., 0.
recalls = torch.tensor(recalls)
precisions = torch.tensor(precisions)
weights = torch.tensor(weights)
total_number = torch.sum(weights)
aver_recall = torch.sum(recalls * weights) / total_number
aver_precision = torch.sum(precisions * weights) / total_number
aver_f1 = 0. if (aver_recall + aver_precision) = 0 else aver_recall * aver_precision / (aver_recall + aver_precision)
return aver_recall, aver_precision, aver_f1
def validate(configs):
# read configs
## command line config
use_gpu = configs['use_gpu']
save_dir = configs['save_dir']
data_root = configs['data_root']
## data cofig
data_config = configs['data']
validation_data_name = data_config['VAL']
## superpoint model config
superpoint_model_config = configs['model']['superpoint']
detection_threshold = superpoint_model_config['eval']['detection_threshold']
## graph model config
gcn_config = configs['model']['gcn']
batch_szie = gcn_config['train']['batch_szie']
## others
configs['num_gpu'] = [0]
configs['public_model'] = 0
# data
data_loader = coco_loader(data_root=data_root, name=validation_data_name, config=data_config,
batch_size=batch_szie, remove_images_without_annotations=True)
# model
maskrcnn_model = build_maskrcnn(configs)
gcn_model = build_gcn(configs)
recall, precision, f1 = calculate_f1(maskrcnn_model, gcn_model, data_loader, configs)
print("recall = {}, precision = {}, f1 = {}".format(recall, precision, f1))
def main():
parser = argparse.ArgumentParser(description="Validation")
parser.add_argument(
"-c", "--config_file",
dest = "config_file",
type = str,
default = ""
)
parser.add_argument(
"-g", "--gpu",
dest = "gpu",
type = int,
default = 0
)
parser.add_argument(
"-s", "--save_dir",
dest = "save_dir",
type = str,
default = ""
)
parser.add_argument(
"-d", "--data_root",
dest = "data_root",
type = str,
default = ""
)
parser.add_argument(
"-mm", "--maskrcnn_model_path",
dest = "maskrcnn_model_path",
type = str,
default = ""
)
parser.add_argument(
"-gm", "--graph_model_path",
dest = "graph_model_path",
type = str,
default = ""
)
args = parser.parse_args()
config_file = args.config_file
f = open(config_file, 'r', encoding='utf-8')
configs = f.read()
configs = yaml.load(configs)
configs['use_gpu'] = args.gpu
configs['save_dir'] = args.save_dir
configs['data_root'] = args.data_root
configs['pretrained_model_path'] = args.maskrcnn_model_path
configs['graph_model_path'] = args.graph_model_path
validate(configs)
if __name__ == "__main__":
main()