def train_forward(self, training_data):
target_img = training_data["im_z"]
search_img = training_data["im_x"]
# backbone feature
f_z = self.basemodel(target_img)
f_x = self.basemodel(search_img)
# feature adjustment
c_z_k = self.c_z_k(f_z)
r_z_k = self.r_z_k(f_z)
c_x = self.c_x(f_x)
r_x = self.r_x(f_x)
# feature matching
r_out = xcorr_depthwise(r_x, r_z_k)
c_out = xcorr_depthwise(c_x, c_z_k)
# head
fcos_cls_score_final, fcos_ctr_score_final, fcos_bbox_final, corr_fea = self.head(
c_out, r_out)
predict_data = dict(
cls_pred=fcos_cls_score_final,
ctr_pred=fcos_ctr_score_final,
box_pred=fcos_bbox_final,
)
if self._hyper_params["corr_fea_output"]:
predict_data["corr_fea"] = corr_fea
return predict_data
def update(self, im_z, search_img):
# self.target_img.require_grad = True
f_z = self.basemodel(im_z)
# template as kernel
c_z_k = self.c_z_k(f_z)
r_z_k = self.r_z_k(f_z)
# backbone feature
f_x = self.basemodel(search_img)
# feature adjustment
c_x = self.c_x(f_x)
r_x = self.r_x(f_x)
# feature matching
r_out = xcorr_depthwise(r_x, r_z_k)
c_out = xcorr_depthwise(c_x, c_z_k)
# head
fcos_cls_score_final, fcos_ctr_score_final, fcos_bbox_final, corr_fea = self.head(
c_out, r_out)
return fcos_cls_score_final, fcos_ctr_score_final, fcos_bbox_final
def forward(self, *args, phase="train"):
r"""
Perform tracking process for different phases (e.g. train / init / track)
Arguments
---------
target_img: torch.Tensor
target template image patch
search_img: torch.Tensor
search region image patch
Returns
-------
fcos_score_final: torch.Tensor
predicted score for bboxes, shape=(B, HW, 1)
fcos_bbox_final: torch.Tensor
predicted bbox in the crop, shape=(B, HW, 4)
fcos_cls_prob_final: torch.Tensor
classification score, shape=(B, HW, 1)
fcos_ctr_prob_final: torch.Tensor
center-ness score, shape=(B, HW, 1)
"""
if phase == 'train':
# resolve training data
training_data = args[0]
target_img = training_data["im_z"]
search_img = training_data["im_x"]
# backbone feature
f_z = self.basemodel(target_img)
f_x = self.basemodel(search_img)
# feature adjustment
c_z_k = self.c_z_k(f_z)
r_z_k = self.r_z_k(f_z)
c_x = self.c_x(f_x)
r_x = self.r_x(f_x)
# feature matching
r_out = xcorr_depthwise(r_x, r_z_k)
c_out = xcorr_depthwise(c_x, c_z_k)
# head
fcos_cls_score_final, fcos_ctr_score_final, fcos_bbox_final, corr_fea = self.head(
c_out, r_out)
predict_data = dict(
cls_pred=fcos_cls_score_final,
ctr_pred=fcos_ctr_score_final,
box_pred=fcos_bbox_final,
)
if self._hyper_params["corr_fea_output"]:
predict_data["corr_fea"] = corr_fea
return predict_data
elif phase == 'feature':
target_img, = args
# backbone feature
f_z = self.basemodel(target_img)
# template as kernel
c_z_k = self.c_z_k(f_z)
r_z_k = self.r_z_k(f_z)
# output
out_list = [c_z_k, r_z_k]
elif phase == 'track':
if len(args) == 3:
search_img, c_z_k, r_z_k = args
# backbone feature
f_x = self.basemodel(search_img)
# feature adjustment
c_x = self.c_x(f_x)
r_x = self.r_x(f_x)
elif len(args) == 4:
# c_x, r_x already computed
c_z_k, r_z_k, c_x, r_x = args
else:
raise ValueError("Illegal args length: %d" % len(args))
# feature matching
r_out = xcorr_depthwise(r_x, r_z_k)
c_out = xcorr_depthwise(c_x, c_z_k)
# head
fcos_cls_score_final, fcos_ctr_score_final, fcos_bbox_final, corr_fea = self.head(
c_out, r_out, search_img.size(-1))
# apply sigmoid
fcos_cls_prob_final = torch.sigmoid(fcos_cls_score_final)
fcos_ctr_prob_final = torch.sigmoid(fcos_ctr_score_final)
# apply centerness correction
fcos_score_final = fcos_cls_prob_final * fcos_ctr_prob_final
# register extra output
extra = dict(c_x=c_x, r_x=r_x, corr_fea=corr_fea)
# output
out_list = fcos_score_final, fcos_bbox_final, fcos_cls_prob_final, fcos_ctr_prob_final, extra
else:
raise ValueError("Phase non-implemented.")
return out_list
def forward(self, *args, phase=None):
r"""
Perform tracking process for different phases (e.g. train / init / track)
Arguments
---------
target_img: torch.Tensor
target template image patch
search_img: torch.Tensor
search region image patch
Returns
-------
fcos_score_final: torch.Tensor
predicted score for bboxes, shape=(B, HW, 1)
fcos_bbox_final: torch.Tensor
predicted bbox in the crop, shape=(B, HW, 4)
fcos_cls_prob_final: torch.Tensor
classification score, shape=(B, HW, 1)
fcos_ctr_prob_final: torch.Tensor
center-ness score, shape=(B, HW, 1)
"""
if phase is None:
phase = self._phase
# used during training
if phase == 'train':
# resolve training data
if self._hyper_params["amp"]:
with torch.cuda.amp.autocast():
return self.train_forward(args[0])
else:
return self.train_forward(args[0])
# used for template feature extraction (normal mode)
elif phase == 'feature':
target_img, = args
if self._hyper_params["trt_mode"]:
# extract feature with trt model
out_list = self.trt_fea_model(target_img)
else:
# backbone feature
f_z = self.basemodel(target_img)
# template as kernel
c_z_k = self.c_z_k(f_z)
r_z_k = self.r_z_k(f_z)
# output
out_list = [c_z_k, r_z_k]
# used for template feature extraction (trt mode)
elif phase == "freeze_track_fea":
search_img, = args
# backbone feature
f_x = self.basemodel(search_img)
# feature adjustment
c_x = self.c_x(f_x)
r_x = self.r_x(f_x)
# head
return [c_x, r_x]
# [Broken] used for template feature extraction (trt mode)
# currently broken due to following issue of "torch2trt" package
# c.f. https://github.com/NVIDIA-AI-IOT/torch2trt/issues/251
elif phase == "freeze_track_head":
c_out, r_out = args
# head
outputs = self.head(c_out, r_out, 0, True)
return outputs
# used for tracking one frame during test
elif phase == 'track':
if len(args) == 3:
search_img, c_z_k, r_z_k = args
if self._hyper_params["trt_mode"]:
c_x, r_x = self.trt_track_model(search_img)
else:
# backbone feature
f_x = self.basemodel(search_img)
# feature adjustment
c_x = self.c_x(f_x)
r_x = self.r_x(f_x)
elif len(args) == 4:
# c_x, r_x already computed
c_z_k, r_z_k, c_x, r_x = args
else:
raise ValueError("Illegal args length: %d" % len(args))
# feature matching
r_out = xcorr_depthwise(r_x, r_z_k)
c_out = xcorr_depthwise(c_x, c_z_k)
# head
fcos_cls_score_final, fcos_ctr_score_final, fcos_bbox_final, corr_fea = self.head(
c_out, r_out, search_img.size(-1))
# apply sigmoid
fcos_cls_prob_final = torch.sigmoid(fcos_cls_score_final)
fcos_ctr_prob_final = torch.sigmoid(fcos_ctr_score_final)
# apply centerness correction
fcos_score_final = fcos_cls_prob_final * fcos_ctr_prob_final
# register extra output
extra = dict(c_x=c_x, r_x=r_x, corr_fea=corr_fea)
# output
out_list = fcos_score_final, fcos_bbox_final, fcos_cls_prob_final, fcos_ctr_prob_final, extra
else:
raise ValueError("Phase non-implemented.")
return out_list
def forward(self, *args, phase="train"):
r"""
Perform tracking process for different phases (e.g. train / init / track)
Arguments
---------
target_img: torch.Tensor
target template image patch
search_img: torch.Tensor
search region image patch
Returns
-------
fcos_score_final: torch.Tensor
predicted score for bboxes, shape=(B, HW, 1)
fcos_bbox_final: torch.Tensor
predicted bbox in the crop, shape=(B, HW, 4)
fcos_cls_prob_final: torch.Tensor
classification score, shape=(B, HW, 1)
fcos_ctr_prob_final: torch.Tensor
center-ness score, shape=(B, HW, 1)
"""
# phase: train
if phase == 'train':
pass
# phase: feature
elif phase == 'feature':
target_img, = args
# backbone feature
f_z = self.basemodel_target(target_img)
r_z_k = self.r_z_k(f_z)
out_list = [f_z, r_z_k]
# phase: track
elif phase == 'track':
search_img, f_z = args
c_z_k = self.c_z_k(f_z)
r_z_k = self.r_z_k(f_z)
f_x = self.basemodel_search(search_img)
# feature adjustment
c_x = self.c_x(f_x)
r_x = self.r_x(f_x)
# feature matching
r_out = xcorr_depthwise(r_x, r_z_k)
c_out = xcorr_depthwise(c_x, c_z_k)
# head
fcos_cls_score_final, fcos_ctr_score_final, fcos_bbox_final, corr_feature = self.head(
c_out, r_out)
# apply sigmoid
fcos_cls_prob_final = torch.sigmoid(fcos_cls_score_final)
fcos_ctr_prob_final = torch.sigmoid(fcos_ctr_score_final)
# apply centerness correction
fcos_score_final = fcos_cls_prob_final * fcos_ctr_prob_final
# register extra output
#extra = dict(c_x=c_x, r_x=r_x, corr_feature = corr_feature)
# output
out_list = fcos_score_final, fcos_bbox_final, fcos_cls_prob_final, fcos_ctr_prob_final, corr_feature
else:
raise ValueError("Phase non-implemented.")
return out_list
