#setup optimizer
params = list(net.parameters())
# optimizer = torch.optim.Adam(params[-8:], lr=lr)
#optimizer = torch.optim.SGD(params[8:], lr=lr, momentum=momentum, weight_decay=weight_decay)
optimizer = torch.optim.SGD(params, lr=lr, momentum=momentum, weight_decay=weight_decay)
#make sure dir for saving model checkpoints exists
if not os.path.exists(output_dir):
os.mkdir(output_dir)
# things to print out during training training
train_loss = 0
step_cnt = 0
t = Timer()
t.tic()
for epoch in range(num_epochs):
#more things to print out later
tv_cnt = 0
ir_cnt = 0
targets_cnt = {}#how many times a target is used(visible, total)
for cid in chosen_ids:
targets_cnt[cid] = [0,0]
epoch_loss = 0
epoch_step_cnt = 0
for step,batch in enumerate(trainloader):
step = step+1
def test_net(
model_name,
net,
dataset,
num_images=100,
batch=False,
max_per_target=5,
thresh=0.05,
vis=False,
output_dir=None,
):
"""Test a network on an image database."""
_t = {'im_detect': Timer(), 'misc': Timer()}
correct = 0
correct_ish = 0
total = 0
for il in range(num_images):
#gets a random batch
batch = dataset[0]
#get first and second images
im_data = match_and_concat_images(batch[0], batch[3]) - 127
#get target images
target_data = match_and_concat_images(batch[2][0], batch[2][1]) - 127
#get gt box, add 1 for fg label
gt_box = batch[1]
gt_box.append(1)
gt_box = np.asarray(gt_box, dtype=np.float32)
#1 gt_box for each image, the second is a dummy box since there is no fg
gt_boxes = np.asarray([gt_box, [0, 0, 1, 1, 0]])
_t['im_detect'].tic()
scores, boxes = im_detect(net, target_data, im_data)
detect_time = _t['im_detect'].toc(average=False)
_t['misc'].tic()
#get scores for foreground, non maximum supression
inds = np.where(scores[:] > thresh)[0]
#fg_scores = scores[inds, 1]
fg_scores = scores[inds, :]
fg_boxes = boxes[inds, 1 * 4:(1 + 1) * 4]
#fg_dets = np.hstack((fg_boxes, fg_scores[:, np.newaxis])) \
# .astype(np.float32, copy=False)
fg_dets = np.hstack((fg_boxes, fg_scores)) \
.astype(np.float32, copy=False)
keep = nms(fg_dets, cfg.TEST.NMS)
fg_dets = fg_dets[keep, :]
# Limit to max_per_target detections *over all classes*
if max_per_target > 0:
image_scores = np.hstack([fg_dets[:, -1]])
if len(image_scores) > max_per_target:
image_thresh = np.sort(image_scores)[-max_per_target]
keep = np.where(fg_dets[:, -1] >= image_thresh)[0]
fg_dets = fg_dets[keep, :]
if len(fg_dets) > max_per_target:
breakp = 1
nms_time = _t['misc'].toc(average=False)
print 'im_detect: {:d}/{:d} {:.3f}s {:.3f}s' \
.format(il + 1, num_images, detect_time, nms_time)
not_found = True
for box in fg_dets:
iou = get_boxes_iou(gt_box, box)
if not_found and iou > .5 and box[-1] > .3:
not_found = False
correct += 1
if iou > .4:
#not_found = False
correct_ish += 1
total += 1
#see if box is correct or not
print correct
print correct_ish
print total
acc = 0
acc_ish = 0
if total != 0:
acc = float(correct) / float(total)
acc_ish = float(correct_ish) / float(total)
return [acc, acc_ish, correct, correct_ish]
def test_net(
model_name,
net,
dataloader,
name_to_id,
target_images,
chosen_ids,
max_per_target=5,
thresh=0.05,
vis=False,
output_dir=None,
):
"""Test a Fast R-CNN network on an image database."""
#get map from target id to name
id_to_name = {}
for name in name_to_id.keys():
id_to_name[name_to_id[name]] = name
#num images in test set
num_images = len(dataloader)
# all detections are collected into:
# all_boxes[cls][image] = N x 5 array of detections in
# (x1, y1, x2, y2, score)
all_boxes = [[[] for _ in xrange(num_images)]
for _ in xrange(dataloader.dataset.get_num_classes())]
#array of result dicts
all_results = {}
# timers
_t = {'im_detect': Timer(), 'misc': Timer()}
if output_dir is not None:
det_file = os.path.join(output_dir, model_name + '.json')
print det_file
#for i in range(num_images):
for i, batch in enumerate(dataloader):
im_data = batch[0].unsqueeze(0).numpy()
im_data = np.transpose(im_data, (0, 2, 3, 1))
im_info = np.zeros((1, 3))
im_info[0, :] = [im_data.shape[1], im_data.shape[2], 1]
dontcare_areas = np.zeros((0, 4))
all_image_dets = np.zeros((0, 6))
for id_ind, t_id in enumerate(chosen_ids):
target_name = id_to_name[t_id]
if target_name == 'background':
continue
target_data = target_images[target_name]
if (target_data is None) or len(target_data) < 1:
print 'Empty target data: {}'.format(target_name)
continue
_t['im_detect'].tic()
scores, boxes = im_detect(net, target_data, im_data, im_info)
detect_time = _t['im_detect'].toc(average=False)
_t['misc'].tic()
#get scores for foreground, non maximum supression
inds = np.where(scores[:, 1] > thresh)[0]
fg_scores = scores[inds, 1]
fg_boxes = boxes[inds, 1 * 4:(1 + 1) * 4]
fg_dets = np.hstack((fg_boxes, fg_scores[:, np.newaxis])) \
.astype(np.float32, copy=False)
keep = nms(fg_dets, cfg.TEST.NMS)
fg_dets = fg_dets[keep, :]
# Limit to max_per_target detections *over all classes*
if max_per_target > 0:
image_scores = np.hstack([fg_dets[:, -1]])
if len(image_scores) > max_per_target:
image_thresh = np.sort(image_scores)[-max_per_target]
keep = np.where(fg_dets[:, -1] >= image_thresh)[0]
fg_dets = fg_dets[keep, :]
nms_time = _t['misc'].toc(average=False)
print 'im_detect: {:d}/{:d} {:.3f}s {:.3f}s' \
.format(i + 1, num_images, detect_time, nms_time)
#put class id in the box
fg_dets = np.insert(fg_dets, 4, t_id, axis=1)
all_image_dets = np.vstack((all_image_dets, fg_dets))
#record results by image name
all_results[batch[1][1]] = all_image_dets.tolist()
if output_dir is not None:
with open(det_file, 'w') as f:
json.dump(all_results, f)
return all_results