-
Notifications
You must be signed in to change notification settings - Fork 182
/
train.py
131 lines (117 loc) · 5.72 KB
/
train.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
from __future__ import division
import os
import numpy as np
from torch.autograd import Variable
from torch.utils import data as data_
import torch
import random
import ipdb
import argparse
import cv2
from src.head_detector_vgg16 import Head_Detector_VGG16
from trainer import Head_Detector_Trainer
from src.config import opt
import src.utils as utils
from data.dataset import Dataset, inverse_normalize
import src.array_tool as at
from src.vis_tool import visdom_bbox
from src.bbox_tools import bbox_iou
dataset_name = 'hollywood'
phases = ['train', 'val', 'test']
data_check_flag = False
def eval(dataloader, head_detector):
"""
Given the dataloader of the test split compute the
average corLoc of the dataset using the head detector
model given as the argument to the function.
"""
test_img_num = 0
test_corrLoc = 0.0
for _, (img, bbox_, scale) in enumerate(dataloader):
scale = at.scalar(scale)
img, bbox = img.cuda().float(), bbox_.cuda()
img, bbox = Variable(img), Variable(bbox)
pred_bboxes_, _ = head_detector.predict(img, scale, mode='evaluate')
gt_bboxs = at.tonumpy(bbox_)[0]
pred_bboxes_ = at.tonumpy(pred_bboxes_)
if pred_bboxes_.shape[0] == 0:
test_img_num += 1
continue
else:
ious = bbox_iou(pred_bboxes_, gt_bboxs)
max_ious = ious.max(axis=1)
corr_preds = np.where(max_ious >= 0.5)[0]
num_boxs = gt_bboxs.shape[0]
num_corr_preds = len(corr_preds)
test_corrLoc += num_corr_preds / num_boxs
test_img_num += 1
return test_corrLoc / test_img_num
def train():
# Get the dataset
for phase in phases:
if phase == 'train':
if dataset_name == 'hollywood':
train_data_list_path = os.path.join(opt.hollywood_dataset_root_path, 'hollywood_train.idl')
train_data_list = utils.get_phase_data_list(train_data_list_path, dataset_name)
if dataset_name == 'brainwash':
train_data_list_path = os.path.join(opt.brainwash_dataset_root_path, 'brainwash_train.idl')
train_data_list = utils.get_phase_data_list(train_data_list_path, dataset_name)
elif phase == 'val':
if dataset_name == 'hollywood':
val_data_list_path = os.path.join(opt.hollywood_dataset_root_path, 'hollywood_val.idl')
val_data_list = utils.get_phase_data_list(val_data_list_path, dataset_name)
if dataset_name == 'brainwash':
val_data_list_path = os.path.join(opt.brainwash_dataset_root_path, 'brainwash_val.idl')
val_data_list = utils.get_phase_data_list(val_data_list_path, dataset_name)
elif phase == 'test':
if dataset_name == 'hollywood':
test_data_list_path = os.path.join(opt.hollywood_dataset_root_path, 'hollywood_test.idl')
test_data_list = utils.get_phase_data_list(test_data_list_path, dataset_name)
if dataset_name == 'brainwash':
test_data_list_path = os.path.join(opt.brainwash_dataset_root_path, 'brainwash_test.idl')
test_data_list = utils.get_phase_data_list(test_data_list_path, dataset_name)
print "Number of images for training: %s" %(len(train_data_list))
print "Number of images for val: %s" %(len(val_data_list))
print "Number of images for test: %s" %(len(test_data_list))
if data_check_flag:
utils.check_loaded_data(train_data_list[random.randint(1,len(train_data_list))])
utils.check_loaded_data(val_data_list[random.randint(1,len(val_data_list))])
utils.check_loaded_data(test_data_list[random.randint(1,len(test_data_list))])
# Load the train dataset
train_dataset = Dataset(train_data_list)
test_dataset = Dataset(val_data_list)
print "Load data."
train_dataloader = data_.DataLoader(train_dataset, batch_size=1,shuffle=True, num_workers=1)
test_dataloader = data_.DataLoader(test_dataset, batch_size=1, shuffle=True, num_workers=1)
# Initialize the head detector.
head_detector_vgg16 = Head_Detector_VGG16(ratios=[1], anchor_scales=[8,16])
print "model construct completed"
trainer = Head_Detector_Trainer(head_detector_vgg16).cuda()
lr_ = opt.lr
for epoch in range(opt.epoch):
trainer.reset_meters()
for ii, (img, bbox_, scale) in enumerate(train_dataloader):
scale = at.scalar(scale)
img, bbox = img.cuda().float(), bbox_.cuda()
img, bbox = Variable(img), Variable(bbox)
_, _, _ = trainer.train_step(img, bbox, scale)
print "Forward and backward pass done."
if (ii+1) % opt.plot_every == 0:
trainer.vis.plot_many(trainer.get_meter_data())
ori_img_ = inverse_normalize(at.tonumpy(img[0]))
gt_img = visdom_bbox(ori_img_, at.tonumpy(bbox_[0]))
trainer.vis.img('gt_img', gt_img)
rois,_ = trainer.head_detector.predict(img, scale=scale, mode='visualize')
pred_img = visdom_bbox(ori_img_, at.tonumpy(rois))
trainer.vis.img('pred_img', pred_img)
trainer.vis.text(str(trainer.rpn_cm.value().tolist()), win='rpn_cm')
avg_test_CorrLoc = eval(test_dataloader, head_detector_vgg16)
print("Epoch {} of {}.".format(epoch+1, opt.epoch))
print(" test average corrLoc accuracy:\t\t{:.3f}".format(avg_test_CorrLoc))
model_save_path = trainer.save(best_map=avg_test_CorrLoc)
if epoch == 8:
trainer.load(model_save_path)
trainer.head_detector.scale_lr(opt.lr_decay)
lr_ = lr_ * opt.lr_decay
if __name__ == "__main__":
train()