def main():
opt = parse_args()
use_cuda = cuda_mode(opt)
save_dir = os.path.join(opt.out_dir, opt.dataset_name)
print('Trained models will be save to', os.path.abspath(save_dir))
if not os.path.exists(save_dir):
os.makedirs(save_dir)
pre_model = opt.pre_model
if opt.resume != 0:
pre_model = ''
model = SIPN(opt.net, opt.dataset_name, pre_model)
if opt.resume:
resume = os.path.join(save_dir, 'sipn_' + opt.net + '_' +
str(opt.resume) + '.pth')
print('Resuming model check point from {}'.format(resume))
model.load_trained_model(torch.load(resume))
# Load the dataset
dataset = PersonSearchDataset(opt.data_dir, opt.dataset_name)
# Choose parameters to be updated during training
lr = opt.lr
params = []
print('These parameters will be updated during training:')
for key, value in dict(model.named_parameters()).items():
if value.requires_grad:
print(key)
# TODO: set different decay for weight and bias
params += [{'params': [value], 'lr': lr, 'weight_decay': 1e-4}]
if opt.optimizer == 'SGD':
optimizer = torch.optim.SGD(params, momentum=0.9)
elif opt.optimizer == 'Adam':
# lr *= 0.1
optimizer = torch.optim.Adam(params)
else:
raise KeyError(opt.optimizer)
# Train the model
train_model(dataset, model, lr, optimizer, opt.epochs, use_cuda, save_dir,
opt.resume)
print('Done.\n')
def main():
"""Test the model"""
opt = parse_args()
test_result_dir = os.path.join(opt.out_dir, 'test_result')
dataset_gallery = PersonSearchDataset(opt.data_dir,
split_name='test',
gallery_size=200)
if opt.use_saved_result:
gboxes_file = open(os.path.join(test_result_dir, 'gboxes.pkl'), 'rb')
g_boxes = pickle.load(gboxes_file)
gfeatures_file = open(os.path.join(test_result_dir, 'gfeatures.pkl'),
'rb')
g_features = pickle.load(gfeatures_file)
qfeatures_file = open(os.path.join(test_result_dir, 'qfeatures.pkl'),
'rb')
q_features = pickle.load(qfeatures_file)
else:
use_cuda = cuda_mode(opt)
trained_model_dir = os.path.join('./output',
'sipn_' + opt.trained_epochs + '.pth')
if not os.path.exists(test_result_dir):
os.makedirs(test_result_dir)
net = SIPN(opt.net, trained_model_dir, is_train=False)
net.eval()
if use_cuda:
net.cuda()
# load trained model
print('Loading model check point from {:s}'.format(trained_model_dir))
net.load_trained_model(torch.load(trained_model_dir))
dataset_query = PersonSearchDataset(opt.data_dir,
split_name='test',
gallery_size=200,
test_mode='query')
g_boxes, g_features = test_gallery(net, dataset_gallery, use_cuda,
test_result_dir)
q_features = test_query(net, dataset_query, use_cuda, test_result_dir)
# dataset_gallery.evaluate_detections(g_boxes, det_thresh=0.5)
dataset_gallery.evaluate_search(g_boxes,
g_features,
q_features,
det_thresh=0.5,
gallery_size=200,
dump_json=None)
def main():
opt = parse_args()
use_cuda = cuda_mode(opt)
model = SIPN(opt.net, opt.pre_model)
# Load the dataset
dataset = PersonSearchDataset(opt.data_dir)
save_dir = opt.out_dir
print('Trained models will be save to', os.path.abspath(save_dir))
if not os.path.exists(save_dir):
os.makedirs(save_dir)
# Choose parameters to be updated during training
lr = opt.lr
params = []
print('These parameters will be updated during training:')
for key, value in dict(model.named_parameters()).items():
if value.requires_grad:
print(key)
# TODO: set different decay for weight and bias
params += [{'params': [value], 'lr': lr, 'weight_decay': 1e-4}]
if opt.optimizer == 'SGD':
optimizer = torch.optim.SGD(params, momentum=0.9)
elif opt.optimizer == 'Adam':
lr *= 0.1
optimizer = torch.optim.Adam(params)
else:
raise KeyError(opt.optimizer)
# TODO: add resume
# Train the model
train_model(dataset, model, lr, optimizer, opt.epochs, use_cuda, save_dir)
print('Done')
def main():
opt = parse_args()
use_cuda = cuda_mode(opt)
trained_model_dir = os.path.join('./output', 'sipn_' + opt.net + '_' +
opt.trained_epochs + '.pth')
net = SIPN(opt.net, trained_model_dir, is_train=False)
net.eval()
if use_cuda:
net.cuda()
# load trained model
print('Loading model check point from {:s}'.format(trained_model_dir))
net.load_trained_model(torch.load(trained_model_dir))
test_images = os.listdir(opt.data_dir)
# demo_detection(net, opt.data_dir, test_images, use_cuda)
demo_search(net, opt.data_dir, test_images, use_cuda)
def main():
opt = parse_args()
global device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
save_dir = os.path.join(opt.out_dir, opt.dataset_name)
print('Trained models will be save to', os.path.abspath(save_dir))
if not os.path.exists(save_dir):
os.makedirs(save_dir)
pre_model = opt.pre_model
if opt.resume != 0:
pre_model = ''
model = SIPN(opt.net, opt.dataset_name, pre_model)
model.to(device)
if opt.resume:
resume = os.path.join(save_dir, 'sipn_' + opt.net + '_' +
str(opt.resume) + '.pth')
print('Resuming model check point from {}'.format(resume))
model.load_trained_model(torch.load(resume))
# Read the configuration file
with open('config.yml', 'r') as f:
config = yaml.load(f)
target_size = config['target_size']
max_size = config['max_size']
pixel_means = config['pixel_means']
# Compose transformations for the dataset
transform = sipn_transforms.Compose([
sipn_transforms.RandomHorizontalFlip(),
sipn_transforms.Scale(target_size, max_size),
sipn_transforms.ToTensor(),
sipn_transforms.Normalize(pixel_means)
])
# Load the dataset
dataset = SIPNDataset(opt.data_dir, opt.dataset_name, 'train', transform)
dataloader = DataLoader(dataset, shuffle=True, num_workers=8)
# Choose parameters to be updated during training
lr = opt.lr
params = []
print('These parameters will be updated during training:')
for key, value in dict(model.named_parameters()).items():
if value.requires_grad:
print(key)
# TODO: set different decay for weight and bias
params += [{'params': [value], 'lr': lr, 'weight_decay': 1e-4}]
if opt.optimizer == 'SGD':
optimizer = torch.optim.SGD(params, momentum=0.9)
elif opt.optimizer == 'Adam':
# lr *= 0.1
optimizer = torch.optim.Adam(params)
else:
raise KeyError(opt.optimizer)
# Train the model
train_model(dataloader, model, lr, optimizer, opt.epochs, save_dir,
opt.resume)
print('Done.\n')
color='white')
plt.tight_layout()
fig.savefig(os.path.join(im_dir, 'result_' + im_name))
plt.show()
plt.close(fig)
if __name__ == '__main__':
opt = parse_args()
use_cuda = cuda_mode(opt)
trained_model_dir = os.path.join('./output',
'sipn_' + opt.trained_epochs + '.pth')
net = SIPN(opt.net, trained_model_dir, training=False)
net.eval()
if use_cuda:
net.cuda()
# load trained model
print('Loading model check point from {:s}'.format(trained_model_dir))
net.load_trained_model(torch.load(trained_model_dir))
test_images = os.listdir(opt.data_dir)
# demo_detection(net, opt.data_dir, test_images, use_cuda)
demo_search(net, opt.data_dir, test_images, use_cuda)
print('Done')