class DIMP_ALPHA(object):
def __init__(self, tracker_name='dimp', para_name='dimp50_vot19',
refine_model_name='ARcm_coco_seg', threshold=0.15):
self.THRES = threshold
'''create tracker'''
'''DIMP'''
tracker_info = Tracker(tracker_name, para_name, None)
params = tracker_info.get_parameters()
params.visualization = False
params.debug = False
params.visdom_info = {'use_visdom': False, 'server': '127.0.0.1', 'port': 8097}
self.dimp = tracker_info.tracker_class(params)
'''Alpha-Refine'''
project_path = os.path.join(os.path.dirname(__file__),'..','..')
refine_root = os.path.join(project_path, 'ltr/checkpoints/ltr/ARcm_seg/')
refine_path = os.path.join(refine_root, refine_model_name)
'''2020.4.25 input size: 384x384'''
self.alpha = ARcm_seg(refine_path,input_sz=384)
def initialize(self, img_RGB, mask):
region = rect_from_mask(mask)
self.H, self.W, _ = img_RGB.shape
gt_bbox_np = np.array(region).astype(np.float32)
'''Initialize dimp for specific video'''
gt_bbox_torch = torch.from_numpy(gt_bbox_np)
init_info = {}
init_info['init_bbox'] = gt_bbox_torch
_ = self.dimp.initialize(img_RGB, init_info)
'''initilize refinement module for specific video'''
self.alpha.initialize(img_RGB, np.array(gt_bbox_np))
def track(self, img_RGB):
'''TRACK'''
'''base tracker'''
outputs = self.dimp.track(img_RGB)
pred_bbox = outputs['target_bbox']
'''Step1: Post-Process'''
x1, y1, w, h = pred_bbox
# add boundary and min size limit
x1, y1, x2, y2 = bbox_clip(x1, y1, x1 + w, y1 + h, (self.H, self.W))
w = x2 - x1
h = y2 - y1
new_pos = torch.from_numpy(np.array([y1 + h / 2, x1 + w / 2]).astype(np.float32))
new_target_sz = torch.from_numpy(np.array([h, w]).astype(np.float32))
new_scale = torch.sqrt(new_target_sz.prod() / self.dimp.base_target_sz.prod())
##### update
self.dimp.pos = new_pos.clone()
self.dimp.target_sz = new_target_sz
self.dimp.target_scale = new_scale
bbox_new = [x1,y1,w,h]
'''Step2: Mask report'''
pred_mask, search, search_mask = self.alpha.get_mask(img_RGB, np.array(bbox_new),vis=True)
final_mask = (pred_mask > self.THRES).astype(np.uint8)
search_region = search.astype(np.uint8)
search_mask = (search_mask > self.THRES).astype(np.uint8)
return bbox_new, final_mask, search_region, search_mask
class STARK_ALPHA_SEG(object):
def __init__(self,
tracker_name='stark_s',
para_name='baseline',
refine_model_name='ARcm_coco_seg',
threshold=0.15):
self.THRES = threshold
'''create tracker'''
'''STARK'''
tracker_info = Tracker(tracker_name, para_name, "vot20", None)
params = tracker_info.get_parameters()
params.visualization = False
params.debug = False
params.visdom_info = {
'use_visdom': False,
'server': '127.0.0.1',
'port': 8097
}
self.stark = tracker_info.create_tracker(params)
'''Alpha-Refine'''
project_path = os.path.join(os.path.dirname(__file__), '..', '..')
refine_root = os.path.join(project_path,
'ltr/checkpoints/ltr/ARcm_seg/')
refine_path = os.path.join(refine_root, refine_model_name)
'''2020.4.25 input size: 384x384'''
self.alpha = ARcm_seg(refine_path, input_sz=384)
def initialize(self, img_RGB, mask):
region = rect_from_mask(mask)
self.H, self.W, _ = img_RGB.shape
gt_bbox_np = np.array(region).astype(np.float32)
'''Initialize STARK for specific video'''
init_info = {'init_bbox': list(gt_bbox_np)}
_ = self.stark.initialize(img_RGB, init_info)
'''initilize refinement module for specific video'''
self.alpha.initialize(img_RGB, np.array(gt_bbox_np))
def track(self, img_RGB):
'''TRACK'''
'''base tracker'''
outputs = self.stark.track(img_RGB)
pred_bbox = outputs['target_bbox']
'''Step2: Mask report'''
pred_mask, search, search_mask = self.alpha.get_mask(
img_RGB, np.array(pred_bbox), vis=True)
final_mask = (pred_mask > self.THRES).astype(np.uint8)
search_region = search.astype(np.uint8)
search_mask = (search_mask > self.THRES).astype(np.uint8)
return pred_bbox, final_mask, search_region, search_mask