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