add ProContEXT model for video single object tracking

支持ProContEXT视频单目标跟踪-通用领域模型

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11904797

modelscope/metainfo.py

@@ -326,6 +326,7 @@ class Pipelines(object):
     crowd_counting = 'hrnet-crowd-counting'
     action_detection = 'ResNetC3D-action-detection'
     video_single_object_tracking = 'ostrack-vitb-video-single-object-tracking'
+    video_single_object_tracking_procontext = 'procontext-vitb-video-single-object-tracking'
     video_multi_object_tracking = 'video-multi-object-tracking'
     image_panoptic_segmentation = 'image-panoptic-segmentation'
     image_panoptic_segmentation_easycv = 'image-panoptic-segmentation-easycv'

modelscope/models/cv/video_single_object_tracking/models/layers/attn_blocks.py

@@ -38,6 +38,10 @@ def candidate_elimination(attn: torch.Tensor, tokens: torch.Tensor,
     attn_t = attn[:, :, :lens_t, lens_t:]
 
     if box_mask_z is not None:
+        if not isinstance(box_mask_z, list):
+            box_mask_z = [box_mask_z]
+        box_mask_z_cat = torch.stack(box_mask_z, dim=1)
+        box_mask_z = box_mask_z_cat.flatten(1)
         box_mask_z = box_mask_z.unsqueeze(1).unsqueeze(-1).expand(
             -1, attn_t.shape[1], -1, attn_t.shape[-1])
         attn_t = attn_t[box_mask_z]

modelscope/models/cv/video_single_object_tracking/models/layers/head.py

@@ -55,18 +55,17 @@ class CenterPredictor(
             if p.dim() > 1:
                 nn.init.xavier_uniform_(p)
 
-    def forward(self, x, gt_score_map=None):
+    def forward(self, x, return_score=False):
         """ Forward pass with input x. """
         score_map_ctr, size_map, offset_map = self.get_score_map(x)
 
-        # assert gt_score_map is None
-        if gt_score_map is None:
-            bbox = self.cal_bbox(score_map_ctr, size_map, offset_map)
+        if return_score:
+            bbox, max_score = self.cal_bbox(
+                score_map_ctr, size_map, offset_map, return_score=True)
+            return score_map_ctr, bbox, size_map, offset_map, max_score
         else:
-            bbox = self.cal_bbox(
-                gt_score_map.unsqueeze(1), size_map, offset_map)
-
-        return score_map_ctr, bbox, size_map, offset_map
+            bbox = self.cal_bbox(score_map_ctr, size_map, offset_map)
+            return score_map_ctr, bbox, size_map, offset_map
 
     def cal_bbox(self,
                  score_map_ctr,

modelscope/models/cv/video_single_object_tracking/models/ostrack/ostrack.py

@@ -49,13 +49,13 @@ class OSTrack(nn.Module):
         feat_last = x
         if isinstance(x, list):
             feat_last = x[-1]
-        out = self.forward_head(feat_last, None)
+        out = self.forward_head(feat_last)
 
         out.update(aux_dict)
         out['backbone_feat'] = x
         return out
 
-    def forward_head(self, cat_feature, gt_score_map=None):
+    def forward_head(self, cat_feature):
         """
         cat_feature: output embeddings of the backbone, it can be (HW1+HW2, B, C) or (HW2, B, C)
         """
@@ -67,8 +67,7 @@ class OSTrack(nn.Module):
 
         if self.head_type == 'CENTER':
             # run the center head
-            score_map_ctr, bbox, size_map, offset_map = self.box_head(
-                opt_feat, gt_score_map)
+            score_map_ctr, bbox, size_map, offset_map = self.box_head(opt_feat)
             outputs_coord = bbox
             outputs_coord_new = outputs_coord.view(bs, Nq, 4)
             out = {

modelscope/models/cv/video_single_object_tracking/models/procontext/procontext.py

@@ -0,0 +1,110 @@
+# The ProContEXT implementation is also open-sourced by the authors,
+# and available at https://github.com/jp-lan/ProContEXT
+import torch
+from torch import nn
+
+from modelscope.models.cv.video_single_object_tracking.models.layers.head import \
+    build_box_head
+from .vit_ce import vit_base_patch16_224_ce
+
+
+class ProContEXT(nn.Module):
+    """ This is the base class for ProContEXT """
+
+    def __init__(self,
+                 transformer,
+                 box_head,
+                 aux_loss=False,
+                 head_type='CORNER'):
+        """ Initializes the model.
+        Parameters:
+            transformer: torch module of the transformer architecture.
+            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
+        """
+        super().__init__()
+        self.backbone = transformer
+        self.box_head = box_head
+
+        self.aux_loss = aux_loss
+        self.head_type = head_type
+        if head_type == 'CORNER' or head_type == 'CENTER':
+            self.feat_sz_s = int(box_head.feat_sz)
+            self.feat_len_s = int(box_head.feat_sz**2)
+
+    def forward(
+        self,
+        template: torch.Tensor,
+        search: torch.Tensor,
+        ce_template_mask=None,
+        ce_keep_rate=None,
+    ):
+        x, aux_dict = self.backbone(
+            z=template,
+            x=search,
+            ce_template_mask=ce_template_mask,
+            ce_keep_rate=ce_keep_rate,
+        )
+
+        # Forward head
+        feat_last = x
+        if isinstance(x, list):
+            feat_last = x[-1]
+        out = self.forward_head(feat_last, None)
+
+        out.update(aux_dict)
+        out['backbone_feat'] = x
+        return out
+
+    def forward_head(self, cat_feature, gt_score_map=None):
+        """
+        cat_feature: output embeddings of the backbone, it can be (HW1+HW2, B, C) or (HW2, B, C)
+        """
+        enc_opt = cat_feature[:, -self.
+                              feat_len_s:]  # encoder output for the search region (B, HW, C)
+        opt = (enc_opt.unsqueeze(-1)).permute((0, 3, 2, 1)).contiguous()
+        bs, Nq, C, HW = opt.size()
+        opt_feat = opt.view(-1, C, self.feat_sz_s, self.feat_sz_s)
+
+        if self.head_type == 'CENTER':
+            # run the center head
+            score_map_ctr, bbox, size_map, offset_map, score = self.box_head(
+                opt_feat, return_score=True)
+            outputs_coord = bbox
+            outputs_coord_new = outputs_coord.view(bs, Nq, 4)
+            out = {
+                'pred_boxes': outputs_coord_new,
+                'score_map': score_map_ctr,
+                'size_map': size_map,
+                'offset_map': offset_map,
+                'score': score
+            }
+            return out
+        else:
+            raise NotImplementedError
+
+
+def build_procontext(cfg):
+    if cfg.MODEL.BACKBONE.TYPE == 'vit_base_patch16_224_ce':
+        backbone = vit_base_patch16_224_ce(
+            False,
+            drop_path_rate=cfg.MODEL.BACKBONE.DROP_PATH_RATE,
+            ce_loc=cfg.MODEL.BACKBONE.CE_LOC,
+            ce_keep_ratio=cfg.MODEL.BACKBONE.CE_KEEP_RATIO,
+        )
+        hidden_dim = backbone.embed_dim
+        patch_start_index = 1
+    else:
+        raise NotImplementedError
+
+    backbone.finetune_track(cfg=cfg, patch_start_index=patch_start_index)
+
+    box_head = build_box_head(cfg, hidden_dim)
+
+    model = ProContEXT(
+        backbone,
+        box_head,
+        aux_loss=False,
+        head_type=cfg.MODEL.HEAD.TYPE,
+    )
+
+    return model

modelscope/models/cv/video_single_object_tracking/models/procontext/utils.py

@@ -0,0 +1,22 @@
+# The ProContEXT implementation is also open-sourced by the authors,
+# and available at https://github.com/jp-lan/ProContEXT
+import torch
+
+
+def combine_multi_tokens(template_tokens, search_tokens, mode='direct'):
+    if mode == 'direct':
+        if not isinstance(template_tokens, list):
+            merged_feature = torch.cat((template_tokens, search_tokens), dim=1)
+        elif len(template_tokens) >= 2:
+            merged_feature = torch.cat(
+                (template_tokens[0], template_tokens[1]), dim=1)
+            for i in range(2, len(template_tokens)):
+                merged_feature = torch.cat(
+                    (merged_feature, template_tokens[i]), dim=1)
+            merged_feature = torch.cat((merged_feature, search_tokens), dim=1)
+        else:
+            merged_feature = torch.cat(
+                (template_tokens[0], template_tokens[1]), dim=1)
+    else:
+        raise NotImplementedError
+    return merged_feature

modelscope/models/cv/video_single_object_tracking/models/procontext/vit_ce.py

@@ -0,0 +1,128 @@
+# The ProContEXT implementation is also open-sourced by the authors,
+# and available at https://github.com/jp-lan/ProContEXT
+from functools import partial
+
+import torch
+import torch.nn as nn
+from timm.models.layers import to_2tuple
+
+from modelscope.models.cv.video_single_object_tracking.models.layers.attn_blocks import \
+    CEBlock
+from modelscope.models.cv.video_single_object_tracking.models.layers.patch_embed import \
+    PatchEmbed
+from modelscope.models.cv.video_single_object_tracking.models.ostrack.utils import (
+    combine_tokens, recover_tokens)
+from modelscope.models.cv.video_single_object_tracking.models.ostrack.vit_ce import \
+    VisionTransformerCE
+from .utils import combine_multi_tokens
+
+
+class VisionTransformerCE_ProContEXT(VisionTransformerCE):
+    """ Vision Transformer with candidate elimination (CE) module
+
+    A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`
+        - https://arxiv.org/abs/2010.11929
+
+    Includes distillation token & head support for `DeiT: Data-efficient Image Transformers`
+        - https://arxiv.org/abs/2012.12877
+    """
+
+    def forward_features(
+        self,
+        z,
+        x,
+        mask_x=None,
+        ce_template_mask=None,
+        ce_keep_rate=None,
+    ):
+        B = x.shape[0]
+
+        x = self.patch_embed(x)
+        x += self.pos_embed_x
+        if not isinstance(z, list):
+            z = self.patch_embed(z)
+            z += self.pos_embed_z
+            lens_z = self.pos_embed_z.shape[1]
+            x = combine_tokens(z, x, mode=self.cat_mode)
+        else:
+            z_list = []
+            for zi in z:
+                z_list.append(self.patch_embed(zi) + self.pos_embed_z)
+            lens_z = self.pos_embed_z.shape[1] * len(z_list)
+            x = combine_multi_tokens(z_list, x, mode=self.cat_mode)
+
+        x = self.pos_drop(x)
+
+        lens_x = self.pos_embed_x.shape[1]
+
+        global_index_t = torch.linspace(0, lens_z - 1, lens_z).to(x.device)
+        global_index_t = global_index_t.repeat(B, 1)
+
+        global_index_s = torch.linspace(0, lens_x - 1, lens_x).to(x.device)
+        global_index_s = global_index_s.repeat(B, 1)
+        removed_indexes_s = []
+        for i, blk in enumerate(self.blocks):
+            x, global_index_t, global_index_s, removed_index_s, attn = \
+                blk(x, global_index_t, global_index_s, mask_x, ce_template_mask, ce_keep_rate)
+
+            if self.ce_loc is not None and i in self.ce_loc:
+                removed_indexes_s.append(removed_index_s)
+
+        x = self.norm(x)
+        lens_x_new = global_index_s.shape[1]
+        lens_z_new = global_index_t.shape[1]
+
+        z = x[:, :lens_z_new]
+        x = x[:, lens_z_new:]
+
+        if removed_indexes_s and removed_indexes_s[0] is not None:
+            removed_indexes_cat = torch.cat(removed_indexes_s, dim=1)
+
+            pruned_lens_x = lens_x - lens_x_new
+            pad_x = torch.zeros([B, pruned_lens_x, x.shape[2]],
+                                device=x.device)
+            x = torch.cat([x, pad_x], dim=1)
+            index_all = torch.cat([global_index_s, removed_indexes_cat], dim=1)
+            # recover original token order
+            C = x.shape[-1]
+            x = torch.zeros_like(x).scatter_(
+                dim=1,
+                index=index_all.unsqueeze(-1).expand(B, -1, C).to(torch.int64),
+                src=x)
+
+        x = recover_tokens(x, mode=self.cat_mode)
+
+        # re-concatenate with the template, which may be further used by other modules
+        x = torch.cat([z, x], dim=1)
+
+        aux_dict = {
+            'attn': attn,
+            'removed_indexes_s': removed_indexes_s,  # used for visualization
+        }
+
+        return x, aux_dict
+
+    def forward(self, z, x, ce_template_mask=None, ce_keep_rate=None):
+
+        x, aux_dict = self.forward_features(
+            z,
+            x,
+            ce_template_mask=ce_template_mask,
+            ce_keep_rate=ce_keep_rate,
+        )
+
+        return x, aux_dict
+
+
+def _create_vision_transformer(pretrained=False, **kwargs):
+    model = VisionTransformerCE_ProContEXT(**kwargs)
+    return model
+
+
+def vit_base_patch16_224_ce(pretrained=False, **kwargs):
+    """ ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
+    """
+    model_kwargs = dict(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12, **kwargs)
+    model = _create_vision_transformer(pretrained=pretrained, **model_kwargs)
+    return model

modelscope/models/cv/video_single_object_tracking/tracker/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from .ostrack import OSTrack
+from .procontext import ProContEXT

modelscope/models/cv/video_single_object_tracking/tracker/procontext.py

@@ -0,0 +1,174 @@
+# The ProContEXT implementation is also open-sourced by the authors,
+# and available at https://github.com/jp-lan/ProContEXT
+from copy import deepcopy
+
+import torch
+
+from modelscope.models.cv.video_single_object_tracking.models.procontext.procontext import \
+    build_procontext
+from modelscope.models.cv.video_single_object_tracking.utils.utils import (
+    Preprocessor, clip_box, generate_mask_cond, hann2d, sample_target,
+    transform_image_to_crop)
+
+
+class ProContEXT():
+
+    def __init__(self, ckpt_path, device, cfg):
+        network = build_procontext(cfg)
+        network.load_state_dict(
+            torch.load(ckpt_path, map_location='cpu')['net'], strict=True)
+        self.cfg = cfg
+        if device.type == 'cuda':
+            self.network = network.to(device)
+        else:
+            self.network = network
+        self.network.eval()
+        self.preprocessor = Preprocessor(device)
+        self.state = None
+
+        self.feat_sz = self.cfg.TEST.SEARCH_SIZE // self.cfg.MODEL.BACKBONE.STRIDE
+        # motion constrain
+        if device.type == 'cuda':
+            self.output_window = hann2d(
+                torch.tensor([self.feat_sz, self.feat_sz]).long(),
+                centered=True).to(device)
+        else:
+            self.output_window = hann2d(
+                torch.tensor([self.feat_sz, self.feat_sz]).long(),
+                centered=True)
+        self.frame_id = 0
+        # for save boxes from all queries
+        self.z_dict1 = {}
+        self.z_dict_list = []
+        self.update_intervals = [100]
+
+    def initialize(self, image, info: dict):
+        # crop templates
+        crop_resize_patches = [
+            sample_target(
+                image,
+                info['init_bbox'],
+                factor,
+                output_sz=self.cfg.TEST.TEMPLATE_SIZE)
+            for factor in self.cfg.TEST.TEMPLATE_FACTOR
+        ]
+        z_patch_arr, resize_factor, z_amask_arr = zip(*crop_resize_patches)
+        for idx in range(len(z_patch_arr)):
+            template = self.preprocessor.process(z_patch_arr[idx],
+                                                 z_amask_arr[idx])
+            with torch.no_grad():
+                self.z_dict1 = template
+            self.z_dict_list.append(self.z_dict1)
+        self.box_mask_z = []
+        if self.cfg.MODEL.BACKBONE.CE_LOC:
+            for i in range(len(self.cfg.TEST.TEMPLATE_FACTOR) * 2):
+                template_bbox = self.transform_bbox_to_crop(
+                    info['init_bbox'], resize_factor[0],
+                    template.tensors.device).squeeze(1)
+                self.box_mask_z.append(
+                    generate_mask_cond(self.cfg, 1, template.tensors.device,
+                                       template_bbox))
+
+        # init dynamic templates with static templates
+        for idx in range(len(self.cfg.TEST.TEMPLATE_FACTOR)):
+            self.z_dict_list.append(deepcopy(self.z_dict_list[idx]))
+
+        # save states
+        self.state = info['init_bbox']
+        self.frame_id = 0
+
+    def track(self, image, info: dict = None):
+        H, W, _ = image.shape
+        self.frame_id += 1
+        x_patch_arr, resize_factor, x_amask_arr = sample_target(
+            image,
+            self.state,
+            self.cfg.TEST.SEARCH_FACTOR,
+            output_sz=self.cfg.TEST.SEARCH_SIZE)  # (x1, y1, w, h)
+        search = self.preprocessor.process(x_patch_arr, x_amask_arr)
+
+        with torch.no_grad():
+            x_dict = search
+            # merge the template and the search
+            # run the transformer
+            if isinstance(self.z_dict_list, (list, tuple)):
+                self.z_dict = []
+                for i in range(len(self.cfg.TEST.TEMPLATE_FACTOR) * 2):
+                    self.z_dict.append(self.z_dict_list[i].tensors)
+            out_dict = self.network.forward(
+                template=self.z_dict,
+                search=x_dict.tensors,
+                ce_template_mask=self.box_mask_z)
+
+        # add hann windows
+        pred_score_map = out_dict['score_map']
+        conf_score = out_dict['score']
+        response = self.output_window * pred_score_map
+        pred_boxes = self.network.box_head.cal_bbox(response,
+                                                    out_dict['size_map'],
+                                                    out_dict['offset_map'])
+        pred_boxes = pred_boxes.view(-1, 4)
+        # Baseline: Take the mean of all pred boxes as the final result
+        pred_box = (pred_boxes.mean(dim=0) * self.cfg.TEST.SEARCH_SIZE
+                    / resize_factor).tolist()  # (cx, cy, w, h) [0,1]
+        # get the final box result
+        self.state = clip_box(
+            self.map_box_back(pred_box, resize_factor), H, W, margin=10)
+
+        for idx, update_i in enumerate(self.update_intervals):
+            if self.frame_id % update_i == 0 and conf_score > 0.7:
+                crop_resize_patches2 = [
+                    sample_target(
+                        image,
+                        self.state,
+                        factor,
+                        output_sz=self.cfg.TEST.TEMPLATE_SIZE)
+                    for factor in self.cfg.TEST.TEMPLATE_FACTOR
+                ]
+                z_patch_arr2, _, z_amask_arr2 = zip(*crop_resize_patches2)
+                for idx_s in range(len(z_patch_arr2)):
+                    template_t = self.preprocessor.process(
+                        z_patch_arr2[idx_s], z_amask_arr2[idx_s])
+                    self.z_dict_list[
+                        idx_s
+                        + len(self.cfg.TEST.TEMPLATE_FACTOR)] = template_t
+
+        x1, y1, w, h = self.state
+        x2 = x1 + w
+        y2 = y1 + h
+        return {'target_bbox': [x1, y1, x2, y2]}
+
+    def map_box_back(self, pred_box: list, resize_factor: float):
+        cx_prev, cy_prev = self.state[0] + 0.5 * self.state[2], self.state[
+            1] + 0.5 * self.state[3]
+        cx, cy, w, h = pred_box
+        half_side = 0.5 * self.cfg.TEST.SEARCH_SIZE / resize_factor
+        cx_real = cx + (cx_prev - half_side)
+        cy_real = cy + (cy_prev - half_side)
+        return [cx_real - 0.5 * w, cy_real - 0.5 * h, w, h]
+
+    def transform_bbox_to_crop(self,
+                               box_in,
+                               resize_factor,
+                               device,
+                               box_extract=None,
+                               crop_type='template'):
+        if crop_type == 'template':
+            crop_sz = torch.Tensor(
+                [self.cfg.TEST.TEMPLATE_SIZE, self.cfg.TEST.TEMPLATE_SIZE])
+        elif crop_type == 'search':
+            crop_sz = torch.Tensor(
+                [self.cfg.TEST.SEARCH_SIZE, self.cfg.TEST.SEARCH_SIZE])
+        else:
+            raise NotImplementedError
+
+        box_in = torch.tensor(box_in)
+        if box_extract is None:
+            box_extract = box_in
+        else:
+            box_extract = torch.tensor(box_extract)
+        template_bbox = transform_image_to_crop(
+            box_in, box_extract, resize_factor, crop_sz, normalize=True)
+        template_bbox = template_bbox.view(1, 1, 4).to(device)
+
+        return template_bbox

modelscope/pipelines/cv/video_single_object_tracking_pipeline.py

@@ -7,8 +7,8 @@ import cv2
 from modelscope.metainfo import Pipelines
 from modelscope.models.cv.video_single_object_tracking.config.ostrack import \
     cfg
-from modelscope.models.cv.video_single_object_tracking.tracker.ostrack import \
-    OSTrack
+from modelscope.models.cv.video_single_object_tracking.tracker import (
+    OSTrack, ProContEXT)
 from modelscope.models.cv.video_single_object_tracking.utils.utils import (
     check_box, timestamp_format)
 from modelscope.outputs import OutputKeys
@@ -20,6 +20,9 @@ from modelscope.utils.logger import get_logger
 logger = get_logger()
 
 
+@PIPELINES.register_module(
+    Tasks.video_single_object_tracking,
+    module_name=Pipelines.video_single_object_tracking_procontext)
 @PIPELINES.register_module(
     Tasks.video_single_object_tracking,
     module_name=Pipelines.video_single_object_tracking)
@@ -32,10 +35,14 @@ class VideoSingleObjectTrackingPipeline(Pipeline):
             model: model id on modelscope hub.
         """
         super().__init__(model=model, **kwargs)
-        self.cfg = cfg
         ckpt_path = osp.join(model, ModelFile.TORCH_MODEL_BIN_FILE)
         logger.info(f'loading model from {ckpt_path}')
-        self.tracker = OSTrack(ckpt_path, self.device)
+        if self.cfg.get('tracker', None) == 'ProContEXT':
+            self.tracker = ProContEXT(ckpt_path, self.device, self.cfg)
+        else:
+            self.cfg = cfg
+            self.tracker = OSTrack(ckpt_path, self.device)
+
         logger.info('init tracker done')
 
     def preprocess(self, input) -> Input:

tests/pipelines/test_video_single_object_tracking.py

@@ -14,6 +14,7 @@ class SingleObjectTracking(unittest.TestCase, DemoCompatibilityCheck):
     def setUp(self) -> None:
         self.task = Tasks.video_single_object_tracking
         self.model_id = 'damo/cv_vitb_video-single-object-tracking_ostrack'
+        self.model_id_procontext = 'damo/cv_vitb_video-single-object-tracking_procontext'
 
     @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_run_end2end(self):
@@ -26,6 +27,16 @@ class SingleObjectTracking(unittest.TestCase, DemoCompatibilityCheck):
         show_video_tracking_result(video_path, result[OutputKeys.BOXES],
                                    './tracking_result.avi')
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_run_end2end_procontext(self):
+        video_single_object_tracking = pipeline(
+            Tasks.video_single_object_tracking, model=self.model_id_procontext)
+        video_path = 'data/test/videos/dog.avi'
+        init_bbox = [414, 343, 514, 449]  # [x1, y1, x2, y2]
+        result = video_single_object_tracking((video_path, init_bbox))
+        assert OutputKeys.BOXES in result.keys() and len(
+            result[OutputKeys.BOXES]) == 139
+
     @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
     def test_run_modelhub_default_model(self):
         video_single_object_tracking = pipeline(