def __init__(self, params, model_path=None, **kargs):
super(TrackerSiamRPNBIG, self).__init__(name='SiamRPN',
is_deterministic=True)
self.model = SiameseAlexNet()
self.cuda = torch.cuda.is_available()
self.device = torch.device('cuda:0' if self.cuda else 'cpu')
checkpoint = torch.load(model_path, map_location=self.device)
#print("1")
if 'model' in checkpoint.keys():
self.model.load_state_dict(
torch.load(model_path, map_location=self.device)['model'])
else:
self.model.load_state_dict(
torch.load(model_path, map_location=self.device))
if self.cuda:
self.model = self.model.cuda()
self.model.eval()
self.transforms = transforms.Compose([ToTensor()])
valid_scope = 2 * config.valid_scope + 1
self.anchors = util.generate_anchors(config.total_stride,
config.anchor_base_size,
config.anchor_scales,
config.anchor_ratios, valid_scope)
self.window = np.tile(
np.outer(np.hanning(config.score_size),
np.hanning(config.score_size))[None, :],
[config.anchor_num, 1, 1]).flatten()
self.data_loader = TrackerDataLoader()
def __init__(self, net_path=None, **kargs):
super(TrackerSiamRPN, self).__init__(name='SiamRPN',
is_deterministic=True)
'''setup GPU device if available'''
self.cuda = torch.cuda.is_available()
self.device = torch.device('cuda:0' if self.cuda else 'cpu')
'''setup model'''
self.net = SiameseAlexNet()
if net_path is not None:
self.net.load_state_dict(
torch.load(net_path,
map_location=lambda storage, loc: storage))
if self.cuda:
self.net = self.net.to(self.device)
'''setup optimizer'''
self.optimizer = torch.optim.SGD(self.net.parameters(),
lr=config.lr,
momentum=config.momentum,
weight_decay=config.weight_decay)
self.anchors = util.generate_anchors(config.anchor_total_stride,
config.anchor_base_size,
config.anchor_scales,
config.anchor_ratios,
config.anchor_valid_scope)
def __init__(self, net_path=None, **kargs):
super(TrackerSiamRPN, self).__init__(name='SiamRPN',
is_deterministic=True)
'''setup GPU device if available'''
self.cuda = torch.cuda.is_available()
self.device = torch.device('cuda' if self.cuda else 'cpu')
'''setup model'''
self.net = SiameseAlexNet()
self.net.init_weights()
if net_path is not None:
self.net.load_state_dict(
torch.load(net_path,
map_location=lambda storage, loc: storage))
if self.cuda:
print('Training with GPU.')
self.net = self.net.to(self.device)
'''setup optimizer'''
print("Learning rate: {}".format(config.lr))
self.optimizer = torch.optim.SGD(self.net.parameters(),
lr=config.lr,
momentum=config.momentum,
weight_decay=config.weight_decay)
class TrackerSiamRPN(Tracker):
def __init__(self, net_path=None, **kargs):
super(TrackerSiamRPN, self).__init__(name='SiamRPN',
is_deterministic=True)
'''setup GPU device if available'''
self.cuda = torch.cuda.is_available()
self.device = torch.device('cuda' if self.cuda else 'cpu')
'''setup model'''
self.net = SiameseAlexNet()
self.net.init_weights()
if net_path is not None:
self.net.load_state_dict(
torch.load(net_path,
map_location=lambda storage, loc: storage))
if self.cuda:
print('Training with GPU.')
self.net = self.net.to(self.device)
'''setup optimizer'''
print("Learning rate: {}".format(config.lr))
self.optimizer = torch.optim.SGD(self.net.parameters(),
lr=config.lr,
momentum=config.momentum,
weight_decay=config.weight_decay)
def step(self, epoch, dataset, anchors, i=0, train=True):
if train:
self.net.train()
else:
self.net.eval()
template, detection, regression_target, conf_target = dataset
if self.cuda:
template, detection = template.cuda(), detection.cuda()
regression_target, conf_target = regression_target.cuda(
), conf_target.cuda()
pred_score, pred_regression = self.net(template, detection)
pred_conf = pred_score.reshape(-1, 2, config.size).permute(0, 2, 1)
pred_offset = pred_regression.reshape(-1, 4,
config.size).permute(0, 2, 1)
cls_loss = rpn_cross_entropy_balance(pred_conf,
conf_target,
config.num_pos,
config.num_neg,
anchors,
ohem_pos=config.ohem_pos,
ohem_neg=config.ohem_neg)
reg_loss = rpn_smoothL1(pred_offset,
regression_target,
conf_target,
config.num_pos,
ohem=config.ohem_reg)
loss = cls_loss + config.lamb * reg_loss
# Draw ####################
if i % 100 == 0:
if train:
mode = 'train'
else:
mode = 'val'
anchors_show = anchors
exem_img = template[0].cpu().numpy().transpose(
1, 2, 0)[:, :, ::-1] # (127, 127, 3)
cv2.imwrite('exem_img_{}.png'.format(mode), exem_img)
inst_img = detection[0].cpu().numpy().transpose(
1, 2, 0)[:, :, ::-1] # (255, 255, 3)
topk = 1
cls_pred = F.softmax(pred_conf, dim=2)[0, :, 1]
topk_box = util.get_topk_box(cls_pred,
pred_offset[0],
anchors_show,
topk=topk)
img_box = util.add_box_img(inst_img, topk_box,
color=(0, 0, 255)) # Red
cls_pred = conf_target[0]
gt_box = util.get_topk_box(cls_pred, regression_target[0],
anchors_show)
img_box = util.add_box_img(img_box,
gt_box,
color=(255, 0, 0),
x=1,
y=1) # Blue
cv2.imwrite('pred_inst_gt_{}.png'.format(mode), img_box)
# Draw End ####################
if train:
self.optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(self.net.parameters(), config.clip)
self.optimizer.step()
return cls_loss, reg_loss, loss
'''save model'''
def save(self, model, exp_name_dir, epoch, max_save=10):
util.adjust_learning_rate(self.optimizer, config.gamma)
model_save_dir_pth = '{}/model'.format(exp_name_dir)
if not os.path.exists(model_save_dir_pth):
os.makedirs(model_save_dir_pth)
net_path = os.path.join(model_save_dir_pth,
'model_e%d.pth' % (epoch + 1))
torch.save(model.net.state_dict(), net_path)
if epoch + 1 > max_save:
lastsave = os.path.join(model_save_dir_pth,
'model_e%d.pth' % (epoch + 1 - max_save))
if os.path.exists(lastsave):
os.remove(lastsave)
def adjust_lr(self, epoch):
util.adjust_learning_rate(self.optimizer, config.gamma**epoch)
class TrackerSiamRPNBIG(Tracker):
def __init__(self, params, model_path=None, **kargs):
super(TrackerSiamRPNBIG, self).__init__(name='SiamRPN',
is_deterministic=True)
self.model = SiameseAlexNet()
self.cuda = torch.cuda.is_available()
self.device = torch.device('cuda:0' if self.cuda else 'cpu')
checkpoint = torch.load(model_path, map_location=self.device)
#print("1")
if 'model' in checkpoint.keys():
self.model.load_state_dict(
torch.load(model_path, map_location=self.device)['model'])
else:
self.model.load_state_dict(
torch.load(model_path, map_location=self.device))
if self.cuda:
self.model = self.model.cuda()
self.model.eval()
self.transforms = transforms.Compose([ToTensor()])
valid_scope = 2 * config.valid_scope + 1
self.anchors = util.generate_anchors(config.total_stride,
config.anchor_base_size,
config.anchor_scales,
config.anchor_ratios, valid_scope)
self.window = np.tile(
np.outer(np.hanning(config.score_size),
np.hanning(config.score_size))[None, :],
[config.anchor_num, 1, 1]).flatten()
self.data_loader = TrackerDataLoader()
def _cosine_window(self, size):
"""
get the cosine window
"""
cos_window = np.hanning(int(size[0]))[:, np.newaxis].dot(
np.hanning(int(size[1]))[np.newaxis, :])
cos_window = cos_window.astype(np.float32)
cos_window /= np.sum(cos_window)
return cos_window
def init(self, frame, bbox):
""" initialize siamfc tracker
Args:
frame: an RGB image
bbox: one-based bounding box [x, y, width, height]
"""
frame = np.asarray(frame)
bbox[0] = bbox[0] + bbox[2] / 2
bbox[1] = bbox[1] + bbox[3] / 2
self.pos = np.array(
[bbox[0] + bbox[2] / 2 - 1 / 2,
bbox[1] + bbox[3] / 2 - 1 / 2]) # center x, center y, zero based
self.target_sz = np.array([bbox[2], bbox[3]]) # width, height
self.bbox = np.array([
bbox[0] + bbox[2] / 2 - 1 / 2, bbox[1] + bbox[3] / 2 - 1 / 2,
bbox[2], bbox[3]
])
self.origin_target_sz = np.array([bbox[2], bbox[3]])
# get exemplar img
self.img_mean = np.mean(frame, axis=(0, 1))
exemplar_img, _, _ = self.data_loader.get_exemplar_image(
frame, bbox, config.template_img_size, config.context_amount,
self.img_mean)
#cv2.imshow('exemplar_img', exemplar_img)
# get exemplar feature
exemplar_img = self.transforms(exemplar_img)[None, :, :, :]
if self.cuda:
self.model.track_init(exemplar_img.cuda())
else:
self.model.track_init(exemplar_img)
def update(self, frame):
"""track object based on the previous frame
Args:
frame: an RGB image
Returns:
bbox: tuple of 1-based bounding box(xmin, ymin, xmax, ymax)
"""
frame = np.asarray(frame)
instance_img, _, _, scale_x = self.data_loader.get_instance_image(
frame, self.bbox, config.template_img_size,
config.detection_img_size, config.context_amount, self.img_mean)
#cv2.imshow('instance_img', instance_img)
instance_img = self.transforms(instance_img)[None, :, :, :]
if self.cuda:
pred_score, pred_regression = self.model.track(instance_img.cuda())
else:
pred_score, pred_regression = self.model.track(instance_img)
pred_conf = pred_score.reshape(-1, 2, config.size).permute(0, 2, 1)
pred_offset = pred_regression.reshape(-1, 4,
config.size).permute(0, 2, 1)
delta = pred_offset[0].cpu().detach().numpy()
box_pred = util.box_transform_inv(self.anchors, delta)
score_pred = F.softmax(pred_conf, dim=2)[0, :,
1].cpu().detach().numpy()
s_c = util.change(
util.sz(box_pred[:, 2], box_pred[:, 3]) /
(util.sz_wh(self.target_sz * scale_x))) # scale penalty
r_c = util.change((self.target_sz[0] / self.target_sz[1]) /
(box_pred[:, 2] / box_pred[:, 3])) # ratio penalty
penalty = np.exp(-(r_c * s_c - 1.) * config.penalty_k)
pscore = penalty * score_pred
pscore = pscore * (1 - config.window_influence
) + self.window * config.window_influence
best_pscore_id = np.argmax(pscore)
target = box_pred[best_pscore_id, :] / scale_x
lr = penalty[best_pscore_id] * score_pred[
best_pscore_id] * config.lr_box
res_x = np.clip(target[0] + self.pos[0], 0, frame.shape[1])
res_y = np.clip(target[1] + self.pos[1], 0, frame.shape[0])
res_w = np.clip(self.target_sz[0] * (1 - lr) + target[2] * lr,
config.min_scale * self.origin_target_sz[0],
config.max_scale * self.origin_target_sz[0])
res_h = np.clip(self.target_sz[1] * (1 - lr) + target[3] * lr,
config.min_scale * self.origin_target_sz[1],
config.max_scale * self.origin_target_sz[1])
self.pos = np.array([res_x, res_y])
self.target_sz = np.array([res_w, res_h])
bbox = np.array([res_x, res_y, res_w, res_h])
#print('bbox', bbox)
self.bbox = (np.clip(bbox[0], 0, frame.shape[1]).astype(np.float64),
np.clip(bbox[1], 0, frame.shape[0]).astype(np.float64),
np.clip(bbox[2], 10, frame.shape[1]).astype(np.float64),
np.clip(bbox[3], 10, frame.shape[0]).astype(np.float64))
res_x = res_x - res_w / 2 # x -> x1
res_y = res_y - res_h / 2 # y -> y1
bbox = np.array([res_x, res_y, res_w, res_h])
return bbox
class TrackerSiamRPN(Tracker):
def __init__(self, net_path=None, **kargs):
super(TrackerSiamRPN, self).__init__(name='SiamRPN',
is_deterministic=True)
'''setup GPU device if available'''
self.cuda = torch.cuda.is_available()
self.device = torch.device('cuda:0' if self.cuda else 'cpu')
'''setup model'''
self.net = SiameseAlexNet()
if net_path is not None:
self.net.load_state_dict(
torch.load(net_path,
map_location=lambda storage, loc: storage))
if self.cuda:
self.net = self.net.to(self.device)
'''setup optimizer'''
self.optimizer = torch.optim.SGD(self.net.parameters(),
lr=config.lr,
momentum=config.momentum,
weight_decay=config.weight_decay)
self.anchors = util.generate_anchors(config.anchor_total_stride,
config.anchor_base_size,
config.anchor_scales,
config.anchor_ratios,
config.anchor_valid_scope)
def step(self, epoch, dataset, train=True):
if train:
self.net.train()
else:
self.net.eval()
template, detection, regression_target, conf_target = dataset
if self.cuda:
template, detection = template.cuda(), detection.cuda()
regression_target, conf_target = regression_target.cuda(
), conf_target.cuda()
pred_score, pred_regression = self.net(template, detection)
pred_conf = pred_score.reshape(-1, 2, config.size).permute(0, 2, 1)
pred_offset = pred_regression.reshape(-1, 4,
config.size).permute(0, 2, 1)
cls_loss = rpn_cross_entropy_balance(pred_conf,
conf_target,
config.num_pos,
config.num_neg,
self.anchors,
ohem_pos=config.ohem_pos,
ohem_neg=config.ohem_neg)
reg_loss = rpn_smoothL1(pred_offset,
regression_target,
conf_target,
config.num_pos,
ohem=config.ohem_reg)
loss = cls_loss + config.lamb * reg_loss
if train:
self.optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(self.net.parameters(), config.clip)
self.optimizer.step()
return cls_loss, reg_loss, loss
'''save model'''
def save(self, model, exp_name_dir, epoch):
model_save_dir_pth = '{}/model'.format(exp_name_dir)
if not os.path.exists(model_save_dir_pth):
os.makedirs(model_save_dir_pth)
net_path = os.path.join(model_save_dir_pth,
'model_e%d.pth' % (epoch + 1))
torch.save(model.net.state_dict(), net_path)