コード例 #1
0
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
コード例 #2
0
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