-
Notifications
You must be signed in to change notification settings - Fork 0
/
iou_test.py
104 lines (83 loc) · 3.05 KB
/
iou_test.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
import pickle
import torch
from torchvision import transforms
import os
from bounding_box import BoundingBox
from utils import check_overlap, get_iou, load_model
from torch import nn
from predict import predict
from utils import bounding_box_grad
import numpy as np
from pprint import pprint
pickle_path = 'iou/avg_iou.pickle'
with open(pickle_path, 'rb') as read_file:
avg_iou = pickle.load(read_file)
# constants
checkpoint_folder = '/home/user/Models/Experiment-4/All/resnet50'
test_folder_path = '/home/user/Models/Experiment-4/data/test'
result_file = 'iou/test_iou.pickle'
images_text_file = 'data/images.txt'
bounding_box_file = 'data/bounding_boxes.txt'
height = 224
width = 224
num_channels = 3
num_labels = 200
print('Running...')
transform = transforms.Compose([
transforms.Resize((height, width)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
subfolders = sorted(os.listdir(test_folder_path))
bbox = BoundingBox(test_folder_path, images_text_file, bounding_box_file, height, width)
best_equal = None
best_vary = None
best_equal_value = -1.0
best_vary_value = -1.0
normal_method = None
normal_value = None
for key, value in avg_iou.items():
model_name = key.rsplit('.', 1)[0]
train_method, lambda_1, lambda_2 = model_name.split('_')
if lambda_1 == lambda_2 and value > best_equal_value:
best_equal_value = value
best_equal = key
if value > best_vary_value:
best_vary_value = value
best_vary = key
if train_method == 'normal':
normal_value = value
normal_method = key
print(f'normal: {normal_method} : {normal_value}')
print(f'best_equal: {best_equal} : {best_equal_value}')
print(f'best_vary: {best_vary} : {best_vary_value}')
print('Running on test data...')
def calculate_iou(m_name, gp_id):
checkpoint_path = checkpoint_folder + '/' + m_name
criterion = nn.CrossEntropyLoss()
model = load_model(checkpoint_path, num_labels, gp_id)
average_iou = []
for folder_name in subfolders:
label = int(folder_name.split('.')[0]) - 1
target_tensor = torch.tensor([label])
image_names = os.listdir(test_folder_path + '/' + folder_name)
for img_name in image_names:
img_path = test_folder_path + '/' + folder_name + '/' + img_name
x1_gt, y1_gt, x2_gt, y2_gt = bbox.get_bbox_from_path(img_path)
gt = [x1_gt, y1_gt, x2_gt, y2_gt]
grad = predict(model, img_path, transform, criterion, target_tensor, height, width,
num_channels, gp_id)
pred = bounding_box_grad(grad)
overlap = check_overlap(pred, gt)
if overlap:
iou = get_iou(pred, gt)
else:
iou = 0.0
average_iou.append(iou)
return np.average(average_iou)
test_equal = calculate_iou(best_equal, '0')
test_vary = calculate_iou(best_vary, '3')
result = {'test_equal': test_equal, 'test_vary': test_vary}
pprint(result)
with open(result_file, 'wb') as write_file:
pickle.dump(result, write_file)