158 lines
6.1 KiB
Python
158 lines
6.1 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright 2022-2023 The Alibaba Fundamental Vision Team Authors. All rights reserved.
|
|
|
|
import os
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from modelscope.metainfo import Models
|
|
from modelscope.models.base.base_torch_model import TorchModel
|
|
from modelscope.models.builder import MODELS
|
|
from modelscope.utils.config import Config
|
|
from modelscope.utils.constant import ModelFile, Tasks
|
|
from .blocks import (BboxRegressor, Q2VRankerStage1, Q2VRankerStage2,
|
|
V2QRankerStage1, V2QRankerStage2)
|
|
from .swin_transformer import SwinTransformerV2_1D
|
|
|
|
|
|
@MODELS.register_module(
|
|
Tasks.video_temporal_grounding, module_name=Models.soonet)
|
|
class SOONet(TorchModel):
|
|
"""
|
|
The implementation of 'Scanning Only Once: An End-to-end Framework for Fast Temporal Grounding
|
|
in Long Videos'. The model is dynamically initialized with the following parts:
|
|
- q2v_stage1: calculate qv_ctx_score.
|
|
- v2q_stage1: calculate vq_ctx_score.
|
|
- q2v_stage2: calculate qv_ctn_score.
|
|
- v2q_stage2: calculate vq_ctn_score.
|
|
- regressor: predict the offset of bounding box for each candidate anchor.
|
|
"""
|
|
|
|
def __init__(self, model_dir: str, *args, **kwargs):
|
|
"""
|
|
Initialize SOONet Model
|
|
|
|
Args:
|
|
model_dir: model id or path
|
|
"""
|
|
super().__init__()
|
|
config_path = os.path.join(model_dir, ModelFile.CONFIGURATION)
|
|
self.config = Config.from_file(config_path).hyperparams
|
|
nscales = self.config.nscales
|
|
hidden_dim = self.config.hidden_dim
|
|
snippet_length = self.config.snippet_length
|
|
self.enable_stage2 = self.config.enable_stage2
|
|
self.stage2_topk = self.config.stage2_topk
|
|
self.nscales = nscales
|
|
|
|
self.video_encoder = SwinTransformerV2_1D(
|
|
patch_size=snippet_length,
|
|
in_chans=hidden_dim,
|
|
embed_dim=hidden_dim,
|
|
depths=[2] * nscales,
|
|
num_heads=[8] * nscales,
|
|
window_size=[64] * nscales,
|
|
mlp_ratio=2.,
|
|
qkv_bias=True,
|
|
drop_rate=0.,
|
|
attn_drop_rate=0.,
|
|
drop_path_rate=0.1,
|
|
norm_layer=nn.LayerNorm,
|
|
patch_norm=True,
|
|
use_checkpoint=False,
|
|
pretrained_window_sizes=[0] * nscales)
|
|
|
|
self.q2v_stage1 = Q2VRankerStage1(nscales, hidden_dim)
|
|
self.v2q_stage1 = V2QRankerStage1(nscales, hidden_dim)
|
|
if self.enable_stage2:
|
|
self.q2v_stage2 = Q2VRankerStage2(nscales, hidden_dim,
|
|
snippet_length)
|
|
self.v2q_stage2 = V2QRankerStage2(nscales, hidden_dim)
|
|
self.regressor = BboxRegressor(hidden_dim, self.enable_stage2)
|
|
|
|
# Load trained weights
|
|
model_path = os.path.join(model_dir,
|
|
'SOONet_MAD_VIT-B-32_4Scale_10C.pth')
|
|
state_dict = torch.load(model_path, map_location='cpu')['model']
|
|
self.load_state_dict(state_dict, strict=True)
|
|
|
|
def forward(self, **kwargs):
|
|
if self.training:
|
|
return self.forward_train(**kwargs)
|
|
else:
|
|
return self.forward_test(**kwargs)
|
|
|
|
def forward_train(self, **kwargs):
|
|
raise NotImplementedError
|
|
|
|
def forward_test(self,
|
|
query_feats=None,
|
|
video_feats=None,
|
|
start_ts=None,
|
|
end_ts=None,
|
|
scale_boundaries=None,
|
|
**kwargs):
|
|
"""
|
|
Obtain matching scores and bbox bias of the top-k candidate anchors, with
|
|
pre-extracted query features and video features as input.
|
|
|
|
Args:
|
|
query_feats: the pre-extracted text features.
|
|
video_feats: the pre-extracted video features.
|
|
start_ts: the start timestamps of pre-defined multi-scale anchors.
|
|
end_ts: the end timestamps of pre-defined multi-scale anchors.
|
|
scale_boundaries: the begin and end anchor index for each scale in start_ts and end_ts.
|
|
|
|
Returns:
|
|
[final_scores, bbox_bias, starts, ends]
|
|
"""
|
|
sent_feat = query_feats
|
|
ctx_feats = self.video_encoder(video_feats.permute(0, 2, 1))
|
|
qv_ctx_scores = self.q2v_stage1(ctx_feats, sent_feat)
|
|
if self.enable_stage2:
|
|
hit_indices = list()
|
|
starts = list()
|
|
ends = list()
|
|
filtered_ctx_feats = list()
|
|
for i in range(self.nscales):
|
|
_, indices = torch.sort(
|
|
qv_ctx_scores[i], dim=1, descending=True)
|
|
indices, _ = torch.sort(
|
|
torch.LongTensor(
|
|
list(
|
|
set(indices[:, :self.stage2_topk].flatten().cpu().
|
|
numpy().tolist()))))
|
|
indices = indices.to(video_feats.device)
|
|
hit_indices.append(indices)
|
|
|
|
filtered_ctx_feats.append(
|
|
torch.index_select(ctx_feats[i], 1, indices))
|
|
|
|
scale_first = scale_boundaries[i]
|
|
scale_last = scale_boundaries[i + 1]
|
|
|
|
filtered_start = torch.index_select(
|
|
start_ts[scale_first:scale_last], 0, indices)
|
|
filtered_end = torch.index_select(
|
|
end_ts[scale_first:scale_last], 0, indices)
|
|
starts.append(filtered_start)
|
|
ends.append(filtered_end)
|
|
|
|
starts = torch.cat(starts, dim=0)
|
|
ends = torch.cat(ends, dim=0)
|
|
|
|
qv_merge_scores, qv_ctn_scores, ctn_feats = self.q2v_stage2(
|
|
video_feats, sent_feat, hit_indices, qv_ctx_scores)
|
|
ctx_feats = filtered_ctx_feats
|
|
else:
|
|
ctn_feats = None
|
|
qv_merge_scores = qv_ctx_scores
|
|
starts = start_ts
|
|
ends = end_ts
|
|
|
|
bbox_bias = self.regressor(ctx_feats, ctn_feats, sent_feat)
|
|
final_scores = torch.sigmoid(torch.cat(qv_merge_scores, dim=1))
|
|
|
|
return final_scores, bbox_bias, starts, ends
|