From a10e59c8f3473564cd764feee83fb6713e855bcd Mon Sep 17 00:00:00 2001 From: "lanjinpeng.ljp" Date: Thu, 9 Mar 2023 01:12:58 +0800 Subject: [PATCH] add ProContEXT model for video single object tracking MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 支持ProContEXT视频单目标跟踪-通用领域模型 Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11904797 --- modelscope/metainfo.py | 1 + .../models/layers/attn_blocks.py | 4 + .../models/layers/head.py | 15 +- .../models/ostrack/ostrack.py | 7 +- .../models/procontext/__init__.py | 0 .../models/procontext/procontext.py | 110 +++++++++++ .../models/procontext/utils.py | 22 +++ .../models/procontext/vit_ce.py | 128 +++++++++++++ .../tracker/__init__.py | 3 + .../tracker/procontext.py | 174 ++++++++++++++++++ .../video_single_object_tracking_pipeline.py | 15 +- .../test_video_single_object_tracking.py | 11 ++ 12 files changed, 474 insertions(+), 16 deletions(-) create mode 100644 modelscope/models/cv/video_single_object_tracking/models/procontext/__init__.py create mode 100644 modelscope/models/cv/video_single_object_tracking/models/procontext/procontext.py create mode 100644 modelscope/models/cv/video_single_object_tracking/models/procontext/utils.py create mode 100644 modelscope/models/cv/video_single_object_tracking/models/procontext/vit_ce.py create mode 100644 modelscope/models/cv/video_single_object_tracking/tracker/procontext.py diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 17edf12f..d41445e0 100644 --- a/modelscope/metainfo.py +++ b/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' diff --git a/modelscope/models/cv/video_single_object_tracking/models/layers/attn_blocks.py b/modelscope/models/cv/video_single_object_tracking/models/layers/attn_blocks.py index 702c84f1..4eaa40e7 100644 --- a/modelscope/models/cv/video_single_object_tracking/models/layers/attn_blocks.py +++ b/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] diff --git a/modelscope/models/cv/video_single_object_tracking/models/layers/head.py b/modelscope/models/cv/video_single_object_tracking/models/layers/head.py index e0dc7b59..7d296929 100644 --- a/modelscope/models/cv/video_single_object_tracking/models/layers/head.py +++ b/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, diff --git a/modelscope/models/cv/video_single_object_tracking/models/ostrack/ostrack.py b/modelscope/models/cv/video_single_object_tracking/models/ostrack/ostrack.py index 52704a6c..cd560252 100644 --- a/modelscope/models/cv/video_single_object_tracking/models/ostrack/ostrack.py +++ b/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 = { diff --git a/modelscope/models/cv/video_single_object_tracking/models/procontext/__init__.py b/modelscope/models/cv/video_single_object_tracking/models/procontext/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/video_single_object_tracking/models/procontext/procontext.py b/modelscope/models/cv/video_single_object_tracking/models/procontext/procontext.py new file mode 100644 index 00000000..adb18ae4 --- /dev/null +++ b/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 diff --git a/modelscope/models/cv/video_single_object_tracking/models/procontext/utils.py b/modelscope/models/cv/video_single_object_tracking/models/procontext/utils.py new file mode 100644 index 00000000..b29019cf --- /dev/null +++ b/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 diff --git a/modelscope/models/cv/video_single_object_tracking/models/procontext/vit_ce.py b/modelscope/models/cv/video_single_object_tracking/models/procontext/vit_ce.py new file mode 100644 index 00000000..bd580228 --- /dev/null +++ b/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 diff --git a/modelscope/models/cv/video_single_object_tracking/tracker/__init__.py b/modelscope/models/cv/video_single_object_tracking/tracker/__init__.py index e69de29b..82cc97e0 100644 --- a/modelscope/models/cv/video_single_object_tracking/tracker/__init__.py +++ b/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 diff --git a/modelscope/models/cv/video_single_object_tracking/tracker/procontext.py b/modelscope/models/cv/video_single_object_tracking/tracker/procontext.py new file mode 100644 index 00000000..6a8fdfcc --- /dev/null +++ b/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 diff --git a/modelscope/pipelines/cv/video_single_object_tracking_pipeline.py b/modelscope/pipelines/cv/video_single_object_tracking_pipeline.py index 4169def7..89955a53 100644 --- a/modelscope/pipelines/cv/video_single_object_tracking_pipeline.py +++ b/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: diff --git a/tests/pipelines/test_video_single_object_tracking.py b/tests/pipelines/test_video_single_object_tracking.py index 7f3a9226..e75ccbb0 100644 --- a/tests/pipelines/test_video_single_object_tracking.py +++ b/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(