-
Notifications
You must be signed in to change notification settings - Fork 6
/
data_prep.py
141 lines (114 loc) · 5.1 KB
/
data_prep.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
'''
Data preparation
'''
from settings import *
import numpy as np
import pickle
def calc_iou(box_a, box_b):
"""
Calculate the Intersection Over Union of two boxes
Each box specified by upper left corner and lower right corner:
(x1, y1, x2, y2), where 1 denotes upper left corner, 2 denotes lower right corner
Returns IOU value
"""
# Calculate intersection, i.e. area of overlap between the 2 boxes (could be 0)
# http://math.stackexchange.com/a/99576
x_overlap = max(0, min(box_a[2], box_b[2]) - max(box_a[0], box_b[0]))
y_overlap = max(0, min(box_a[3], box_b[3]) - max(box_a[1], box_b[1]))
intersection = x_overlap * y_overlap
# Calculate union
area_box_a = (box_a[2] - box_a[0]) * (box_a[3] - box_a[1])
area_box_b = (box_b[2] - box_b[0]) * (box_b[3] - box_b[1])
union = area_box_a + area_box_b - intersection
iou = intersection / union
return iou
def find_gt_boxes(data_raw, image_file):
"""
Given (global) feature map sizes, and single training example,
find all default boxes that exceed Jaccard overlap threshold
Returns y_true array that flags the matching default boxes with class ID (-1 means nothing there)
"""
# Pre-process ground-truth data
# Convert absolute coordinates to relative coordinates ranging from 0 to 1
# Read the sign class label (note background class label is 0, sign labels are ints >=1)
signs_data = data_raw[image_file]
signs_class = []
signs_box_coords = [] # relative coordinates
for sign_data in signs_data:
# Find class label
sign_class = sign_data['class']
signs_class.append(sign_class)
# Calculate relative coordinates
# (x1, y1, x2, y2), where 1 denotes upper left corner, 2 denotes lower right corner
abs_box_coords = sign_data['box_coords']
scale = np.array([IMG_W, IMG_H, IMG_W, IMG_H])
box_coords = np.array(abs_box_coords) / scale
signs_box_coords.append(box_coords)
# Initialize y_true to all 0s (0 -> background)
y_true_len = 0
for fm_size in FM_SIZES:
y_true_len += fm_size[0] * fm_size[1] * NUM_DEFAULT_BOXES
y_true_conf = np.zeros(y_true_len)
y_true_loc = np.zeros(y_true_len * 4)
# For each GT box, for each feature map, for each feature map cell, for each default box:
# 1) Calculate the Jaccard overlap (IOU) and annotate the class label
# 2) Count how many box matches we got
# 3) If we got a match, calculate normalized box coordinates and updte y_true_loc
match_counter = 0
for i, gt_box_coords in enumerate(signs_box_coords):
y_true_idx = 0
#for fm_idx, fm_size in enumerate(FM_SIZES):
for fm_size in FM_SIZES:
fm_h, fm_w = fm_size # feature map height and width
for row in range(fm_h):
for col in range(fm_w):
for db in DEFAULT_BOXES:
# Calculate relative box coordinates for this default box
x1_offset, y1_offset, x2_offset, y2_offset = db
abs_db_box_coords = np.array([
max(0, col + x1_offset),
max(0, row + y1_offset),
min(fm_w, col+1 + x2_offset),
min(fm_h, row+1 + y2_offset)
])
scale = np.array([fm_w, fm_h, fm_w, fm_h])
db_box_coords = abs_db_box_coords / scale
# Calculate Jaccard overlap (i.e. Intersection Over Union, IOU) of GT box and default box
iou = calc_iou(gt_box_coords, db_box_coords)
# If box matches, i.e. IOU threshold met
if iou >= IOU_THRESH:
# Update y_true_conf to reflect we found a match, and increment match_counter
y_true_conf[y_true_idx] = signs_class[i]
match_counter += 1
# Calculate normalized box coordinates and update y_true_loc
abs_box_center = np.array([col + 0.5, row + 0.5]) # absolute coordinates of center of feature map cell
abs_gt_box_coords = gt_box_coords * scale # absolute ground truth box coordinates (in feature map grid)
norm_box_coords = abs_gt_box_coords - np.concatenate((abs_box_center, abs_box_center))
y_true_loc[y_true_idx*4 : y_true_idx*4 + 4] = norm_box_coords
y_true_idx += 1
return y_true_conf, y_true_loc, match_counter
def do_data_prep(data_raw):
"""
Create the y_true array
data_raw is the dict mapping image_file -> [{'class': class_int, 'box_coords': (x1, y1, x2, y2)}, {...}, ...]
Return a dict {image_file1: {'y_true_conf': y_true_conf, 'y_true_loc': y_true_loc}, image_file2: ...}
"""
# Prepare the data by populating y_true appropriately
data_prep = {}
for image_file in data_raw.keys():
# Find groud-truth boxes based on Jaccard overlap,
# populate y_true_conf (class labels) and y_true_loc (normalized box coordinates)
y_true_conf, y_true_loc, match_counter = find_gt_boxes(data_raw, image_file)
# Only want data points where we have matching default boxes
if match_counter > 0:
data_prep[image_file] = {'y_true_conf': y_true_conf, 'y_true_loc': y_true_loc}
return data_prep
if __name__ == '__main__':
with open('data_raw_%sx%s.p' % (IMG_W, IMG_H), 'rb') as f:
data_raw = pickle.load(f)
print('Preparing data (i.e. matching boxes)')
data_prep = do_data_prep(data_raw)
with open('data_prep_%sx%s.p' % (IMG_W, IMG_H), 'wb') as f:
pickle.dump(data_prep, f)
print('Done. Saved prepared data to data_prep_%sx%s.p' % (IMG_W, IMG_H))
print('Total images with >=1 matching box: %d' % len(data_prep.keys()))