Add files

.gitignore

@@ -0,0 +1,10 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Temporary data
+/checkpoints
+/flagged
+.DS_Store
+._*

README.md

@@ -1,13 +1,11 @@
 ---
 title: R2 Tuning
-emoji: ⚡
-colorFrom: pink
-colorTo: green
+emoji: 🌀
+colorFrom: blue
+colorTo: purple
 sdk: gradio
 sdk_version: 4.36.1
 app_file: app.py
 pinned: false
 license: bsd-3-clause
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,112 @@
+# Copyright (c) Ye Liu. Licensed under the BSD 3-Clause License.
+
+from functools import partial
+
+import clip
+import decord
+import nncore
+import torch
+import gradio as gr
+import matplotlib.pyplot as plt
+import numpy as np
+import torchvision.transforms.functional as F
+from decord import VideoReader
+from nncore.engine import load_checkpoint
+from nncore.nn import build_model
+
+TITLE = '🌀R2-Tuning: Efficient Image-to-Video Transfer Learning for Video Temporal Grounding'  # noqa
+DESCRIPTION = 'R2-Tuning is a parameter- and memory efficient transfer learning method for video temporal grounding. Please find more details in our <a href="https://arxiv.org/abs/2404.00801" target="_blank">Tech Report</a> and <a href="https://github.com/yeliudev/R2-Tuning" target="_blank">GitHub Repo</a>.\n\nUser Guide:\n1. Upload or record a video using web camera.\n2. Input a text query. A good practice is to use a sentence with 5~10 words.\n3. Click "submit" and you\'ll see the moment retrieval and highlight detection results on the right.'  # noqa
+
+CONFIG = 'configs/qvhighlights/r2_tuning_qvhighlights.py'
+WEIGHT = 'https://huggingface.co/yeliudev/R2-Tuning/resolve/main/checkpoints/r2_tuning_qvhighlights-ed516355.pth'  # noqa
+
+
+def convert_time(seconds):
+    minutes, seconds = divmod(round(seconds), 60)
+    return f'{minutes:02d}:{seconds:02d}'
+
+
+def load_video(video_path, cfg):
+    decord.bridge.set_bridge('torch')
+
+    vr = VideoReader(video_path)
+    stride = vr.get_avg_fps() / cfg.data.val.fps
+    fm_idx = [min(round(i), len(vr) - 1) for i in np.arange(0, len(vr), stride).tolist()]
+    video = vr.get_batch(fm_idx).permute(0, 3, 1, 2).float() / 255
+
+    size = 336 if '336px' in cfg.model.arch else 224
+    h, w = video.size(-2), video.size(-1)
+    s = min(h, w)
+    x, y = round((h - s) / 2), round((w - s) / 2)
+    video = video[..., x:x + s, y:y + s]
+    video = F.resize(video, size=(size, size))
+    video = F.normalize(video, (0.481, 0.459, 0.408), (0.269, 0.261, 0.276))
+    video = video.reshape(video.size(0), -1).unsqueeze(0)
+
+    return video
+
+
+def init_model(config, checkpoint):
+    cfg = nncore.Config.from_file(config)
+    cfg.model.init = True
+
+    if checkpoint.startswith('http'):
+        checkpoint = nncore.download(checkpoint, out_dir='checkpoints')
+
+    model = build_model(cfg.model, dist=False).eval()
+    model = load_checkpoint(model, checkpoint, warning=False)
+
+    return model, cfg
+
+
+def main(video, query, model, cfg):
+    if len(query) == 0:
+        raise gr.Error('Text query can not be empty.')
+
+    try:
+        video = load_video(video, cfg)
+    except Exception:
+        raise gr.Error('Failed to load the video.')
+
+    query = clip.tokenize(query, truncate=True)
+
+    device = next(model.parameters()).device
+    data = dict(video=video.to(device), query=query.to(device), fps=[cfg.data.val.fps])
+
+    with torch.inference_mode():
+        pred = model(data)
+
+    mr = pred['_out']['boundary'][:5].cpu().tolist()
+    mr = [[convert_time(p[0]), convert_time(p[1]), round(p[2], 2)] for p in mr]
+
+    hd = pred['_out']['saliency'].cpu()
+    hd = ((hd - hd.min()) / (hd.max() - hd.min())).tolist()
+
+    fig, ax = plt.subplots(figsize=(10, 5.5))
+    ax.plot(range(0, len(hd) * 2, 2), hd)
+
+    ax.set_xlabel('Time (s)', fontsize=15)
+    ax.set_ylabel('Saliency Score', fontsize=15)
+
+    ax.tick_params(labelsize=14)
+    plt.tight_layout(rect=(0.02, 0.02, 0.95, 0.885))
+
+    return mr, fig
+
+
+model, cfg = init_model(CONFIG, WEIGHT)
+main = partial(main, model=model, cfg=cfg)
+
+demo = gr.Interface(
+    fn=main,
+    inputs=[gr.Video(label='Video'),
+            gr.Textbox(label='Text Query')],
+    outputs=[
+        gr.Dataframe(
+            headers=['Start Time', 'End Time', 'Score'], label='Moment Retrieval'),
+        gr.Plot(label='Highlight Detection')
+    ],
+    allow_flagging='auto',
+    title=TITLE,
+    description=DESCRIPTION)
+demo.launch()

configs/_base_/datasets/qvhighlights.py

@@ -0,0 +1,38 @@
+# dataset settings
+data_type = 'Grounding'
+data_root = 'data/qvhighlights/'
+data = dict(
+    train=dict(
+        type='RepeatDataset',
+        times=4,
+        dataset=dict(
+            type=data_type,
+            label_path=data_root + 'qvhighlights_train.jsonl',
+            video_path=data_root + 'frames_224_0.5fps',
+            cache_path=data_root + 'clip_b32_vid_k4',
+            query_path=data_root + 'clip_b32_txt_k4',
+            use_cache=True,
+            min_video_len=5,
+            fps=0.5,
+            unit=2),
+        loader=dict(batch_size=128, num_workers=4, pin_memory=True, shuffle=True)),
+    val=dict(
+        type=data_type,
+        label_path=data_root + 'qvhighlights_val.jsonl',
+        video_path=data_root + 'frames_224_0.5fps',
+        cache_path=data_root + 'clip_b32_vid_k4',
+        query_path=data_root + 'clip_b32_txt_k4',
+        use_cache=True,
+        fps=0.5,
+        unit=2,
+        loader=dict(batch_size=1, num_workers=4, pin_memory=True, shuffle=False)),
+    test=dict(
+        type=data_type,
+        label_path=data_root + 'qvhighlights_test.jsonl',
+        video_path=data_root + 'frames_224_0.5fps',
+        cache_path=data_root + 'clip_b32_vid_k4',
+        query_path=data_root + 'clip_b32_txt_k4',
+        use_cache=True,
+        fps=0.5,
+        unit=2,
+        loader=dict(batch_size=1, num_workers=4, pin_memory=True, shuffle=False)))

configs/_base_/models/model.py

@@ -0,0 +1,44 @@
+_base_ = ['models']
+# model settings
+model = dict(
+    type='R2Tuning',
+    arch='ViT-B/32',
+    init=False,
+    dims=256,
+    strides=(1, 2, 4, 8),
+    buffer_size=1024,
+    max_num_moment=50,
+    adapter_cfg=dict(
+        type='R2Block',
+        k=4,
+        dropout=0.5,
+        use_tef=True,
+        pos_cfg=dict(type='PositionalEncoding', normalize=True, max_len=1024),
+        tem_cfg=dict(
+            type='TransformerDecoderLayer',
+            heads=8,
+            ratio=4,
+            att_dropout=0.0,
+            ffn_dropout=0.0,
+            att_out_dropout=0.0,
+            ffn_out_dropout=0.0,
+            droppath=0.1,
+            pre_norm=False,
+            bias=True,
+            norm_cfg=dict(type='LN'),
+            act_cfg=dict(type='ReLU', inplace=True),
+            order=('cross_att', 'self_att', 'ffn'),
+            att_init_cfg=dict(type='xavier', distribution='uniform'),
+            ffn_init_cfg=dict(type='kaiming'))),
+    pyramid_cfg=dict(type='ConvPyramid'),
+    pooling_cfg=dict(type='AdaPooling'),
+    class_head_cfg=dict(type='ConvHead', kernal_size=3),
+    coord_head_cfg=dict(type='ConvHead', kernal_size=3),
+    loss_cfg=dict(
+        type='BundleLoss',
+        sample_radius=1.5,
+        loss_cls=dict(type='FocalLoss', loss_weight=1.0),
+        loss_reg=dict(type='L1Loss', loss_weight=0.2),
+        loss_sal=dict(type='SampledNCELoss', loss_weight=0.1),
+        loss_video_cal=dict(type='InfoNCELoss', loss_weight=0.1),
+        loss_layer_cal=dict(type='InfoNCELoss', loss_weight=0.1)))

configs/qvhighlights/r2_tuning_qvhighlights.py

@@ -0,0 +1 @@
+_base_ = ['../_base_/models/model.py', '../_base_/datasets/qvhighlights.py']

models/__init__.py

@@ -0,0 +1,6 @@
+from .adapter import R2Block
+from .blocks import AdaPooling, ConvHead, ConvPyramid
+from .loss import BundleLoss
+from .model import R2Tuning
+
+__all__ = ['R2Block', 'AdaPooling', 'ConvHead', 'ConvPyramid', 'BundleLoss', 'R2Tuning']

models/adapter.py

@@ -0,0 +1,99 @@
+# Copyright (c) Ye Liu. Licensed under the BSD 3-Clause License.
+
+import torch
+import torch.nn as nn
+from nncore.nn import MODELS, build_model
+
+
+@MODELS.register()
+class R2Block(nn.Module):
+
+    def __init__(self,
+                 dims,
+                 in_dims,
+                 k=4,
+                 dropout=0.5,
+                 use_tef=True,
+                 pos_cfg=None,
+                 tem_cfg=None):
+        super(R2Block, self).__init__()
+
+        # yapf:disable
+        self.video_map = nn.Sequential(
+            nn.LayerNorm((in_dims[0] + 2) if use_tef else in_dims[0]),
+            nn.Dropout(dropout),
+            nn.Linear((in_dims[0] + 2) if use_tef else in_dims[0], dims),
+            nn.ReLU(inplace=True),
+            nn.LayerNorm(dims),
+            nn.Dropout(dropout),
+            nn.Linear(dims, dims))
+
+        self.query_map = nn.Sequential(
+            nn.LayerNorm(in_dims[1]),
+            nn.Dropout(dropout),
+            nn.Linear(in_dims[1], dims),
+            nn.ReLU(inplace=True),
+            nn.LayerNorm(dims),
+            nn.Dropout(dropout),
+            nn.Linear(dims, dims))
+        # yapf:enable
+
+        if k > 1:
+            self.gate = nn.Parameter(torch.zeros([k - 1]))
+
+        self.v_map = nn.Linear(dims, dims)
+        self.q_map = nn.Linear(dims, dims)
+        self.scale = nn.Parameter(torch.zeros([k]))
+
+        self.pos = build_model(pos_cfg, dims=dims)
+        self.tem = build_model(tem_cfg, dims=dims)
+
+        self.dims = dims
+        self.in_dims = in_dims
+        self.k = k
+        self.dropout = dropout
+        self.use_tef = use_tef
+
+    def forward(self, video_emb, query_emb, video_msk, query_msk):
+        video_emb = video_emb[-self.k:]
+        query_emb = query_emb[-self.k:]
+
+        _, b, t, p, _ = video_emb.size()
+
+        if self.use_tef:
+            tef_s = torch.arange(0, 1, 1 / t, device=video_emb.device)
+            tef_e = tef_s + 1.0 / t
+            tef = torch.stack((tef_s, tef_e), dim=1)
+            tef = tef.unsqueeze(1).unsqueeze(0).unsqueeze(0).repeat(self.k, b, 1, p, 1)
+            video_emb = torch.cat((video_emb, tef[:, :, :video_emb.size(2)]), dim=-1)
+
+        coll_v, coll_q, last = [], [], None
+        for i in range(self.k - 1, -1, -1):
+            v_emb = self.video_map(video_emb[i])  # B * T * P * C
+            q_emb = self.query_map(query_emb[i])  # B * L * C
+
+            coll_v.append(v_emb[:, :, 0])
+            coll_q.append(q_emb)
+
+            v_pool = v_emb.view(b * t, -1, self.dims)  # BT * P * C
+            q_pool = q_emb.repeat_interleave(t, dim=0)  # BT * L * C
+
+            v_pool_map = self.v_map(v_pool)  # BT * P * C
+            q_pool_map = self.q_map(q_pool)  # BT * L * C
+
+            att = torch.bmm(q_pool_map, v_pool_map.transpose(1, 2)) / self.dims**0.5
+            att = att.softmax(-1)  # BT * L * P
+
+            o_pool = torch.bmm(att, v_pool) + q_pool  # BT * L * C
+            o_pool = o_pool.amax(dim=1, keepdim=True)  # BT * 1 * C
+            v_emb = v_pool[:, 0, None] + o_pool * self.scale[i].tanh()
+            v_emb = v_emb.view(b, t, self.dims)  # B * T * C
+
+            if i < self.k - 1:
+                gate = self.gate[i].sigmoid()
+                v_emb = gate * v_emb + (1 - gate) * last
+
+            v_pe = self.pos(v_emb)
+            last = self.tem(v_emb, q_emb, q_pe=v_pe, q_mask=video_msk, k_mask=query_msk)
+
+        return last, q_emb, coll_v, coll_q

models/blocks.py

@@ -0,0 +1,98 @@
+# Copyright (c) Ye Liu. Licensed under the BSD 3-Clause License.
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from nncore.nn import MODELS
+
+
+class Permute(nn.Module):
+
+    def __init__(self):
+        super(Permute, self).__init__()
+
+    def forward(self, x):
+        return x.transpose(-1, -2)
+
+
+@MODELS.register()
+class ConvPyramid(nn.Module):
+
+    def __init__(self, dims, strides):
+        super(ConvPyramid, self).__init__()
+
+        self.blocks = nn.ModuleList()
+        for s in strides:
+            p = int(math.log2(s))
+            if p == 0:
+                layers = nn.ReLU(inplace=True)
+            else:
+                layers = nn.Sequential()
+                conv_cls = nn.Conv1d if p > 0 else nn.ConvTranspose1d
+                for _ in range(abs(p)):
+                    layers.extend([
+                        Permute(),
+                        conv_cls(dims, dims, 2, stride=2),
+                        Permute(),
+                        nn.LayerNorm(dims),
+                        nn.ReLU(inplace=True)
+                    ])
+            self.blocks.append(layers)
+
+        self.strides = strides
+
+    def forward(self, x, mask, return_mask=False):
+        pymid, pymid_msk = [], []
+
+        for s, blk in zip(self.strides, self.blocks):
+            if x.size(1) < s:
+                continue
+
+            pymid.append(blk(x))
+
+            if return_mask:
+                if s > 1:
+                    msk = F.max_pool1d(mask.float(), s, stride=s).long()
+                elif s < 1:
+                    msk = mask.repeat_interleave(int(1 / s), dim=1)
+                else:
+                    msk = mask
+                pymid_msk.append(msk)
+
+        return pymid, pymid_msk
+
+
+@MODELS.register()
+class AdaPooling(nn.Module):
+
+    def __init__(self, dims):
+        super(AdaPooling, self).__init__()
+        self.att = nn.Linear(dims, 1, bias=False)
+
+    def forward(self, x, mask):
+        a = self.att(x) + torch.where(mask.unsqueeze(2) == 1, .0, float('-inf'))
+        a = a.softmax(dim=1)
+        x = torch.matmul(x.transpose(1, 2), a)
+        x = x.squeeze(2).unsqueeze(1)
+        return x
+
+
+@MODELS.register()
+class ConvHead(nn.Module):
+
+    def __init__(self, dims, out_dims, kernal_size=3):
+        super(ConvHead, self).__init__()
+
+        # yapf:disable
+        self.module = nn.Sequential(
+            Permute(),
+            nn.Conv1d(dims, dims, kernal_size, padding=kernal_size // 2),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(dims, out_dims, kernal_size, padding=kernal_size // 2),
+            Permute())
+        # yapf:enable
+
+    def forward(self, x):
+        return self.module(x)

models/generator.py

@@ -0,0 +1,62 @@
+# Copyright (c) Ye Liu. Licensed under the BSD 3-Clause License.
+
+import torch
+import torch.nn as nn
+
+
+class BufferList(nn.Module):
+
+    def __init__(self, buffers):
+        super(BufferList, self).__init__()
+        for i, buffer in enumerate(buffers):
+            self.register_buffer(str(i), buffer, persistent=False)
+
+    def __len__(self):
+        return len(self._buffers)
+
+    def __iter__(self):
+        return iter(self._buffers.values())
+
+
+class PointGenerator(nn.Module):
+
+    def __init__(self, strides, buffer_size, offset=False):
+        super(PointGenerator, self).__init__()
+
+        reg_range, last = [], 0
+        for stride in strides[1:]:
+            reg_range.append((last, stride))
+            last = stride
+        reg_range.append((last, float('inf')))
+
+        self.strides = strides
+        self.reg_range = reg_range
+        self.buffer_size = buffer_size
+        self.offset = offset
+
+        self.buffer = self._cache_points()
+
+    def _cache_points(self):
+        buffer_list = []
+        for stride, reg_range in zip(self.strides, self.reg_range):
+            reg_range = torch.Tensor([reg_range])
+            lv_stride = torch.Tensor([stride])
+            points = torch.arange(0, self.buffer_size, stride)[:, None]
+            if self.offset:
+                points += 0.5 * stride
+            reg_range = reg_range.repeat(points.size(0), 1)
+            lv_stride = lv_stride.repeat(points.size(0), 1)
+            buffer_list.append(torch.cat((points, reg_range, lv_stride), dim=1))
+        buffer = BufferList(buffer_list)
+        return buffer
+
+    def forward(self, pymid):
+        points = []
+        sizes = [p.size(1) for p in pymid] + [0] * (len(self.buffer) - len(pymid))
+        for size, buffer in zip(sizes, self.buffer):
+            if size == 0:
+                continue
+            assert size <= buffer.size(0), 'reached max buffer size'
+            points.append(buffer[:size, :])
+        points = torch.cat(points)
+        return points

models/loss.py

@@ -0,0 +1,210 @@
+# Copyright (c) Ye Liu. Licensed under the BSD 3-Clause License.
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from nncore.nn import LOSSES, Parameter, build_loss
+
+
+@LOSSES.register()
+class SampledNCELoss(nn.Module):
+
+    def __init__(self,
+                 temperature=0.07,
+                 max_scale=100,
+                 learnable=False,
+                 direction=('row', 'col'),
+                 loss_weight=1.0):
+        super(SampledNCELoss, self).__init__()
+
+        scale = torch.Tensor([math.log(1 / temperature)])
+
+        if learnable:
+            self.scale = Parameter(scale)
+        else:
+            self.register_buffer('scale', scale)
+
+        self.temperature = temperature
+        self.max_scale = max_scale
+        self.learnable = learnable
+        self.direction = (direction, ) if isinstance(direction, str) else direction
+        self.loss_weight = loss_weight
+
+    def extra_repr(self):
+        return ('temperature={}, max_scale={}, learnable={}, direction={}, loss_weight={}'
+                .format(self.temperature, self.max_scale, self.learnable, self.direction,
+                        self.loss_weight))
+
+    def forward(self, video_emb, query_emb, video_msk, saliency, pos_clip):
+        batch_inds = torch.arange(video_emb.size(0), device=video_emb.device)
+
+        pos_scores = saliency[batch_inds, pos_clip].unsqueeze(-1)
+        loss_msk = (saliency <= pos_scores) * video_msk
+
+        scale = self.scale.exp().clamp(max=self.max_scale)
+        i_sim = F.cosine_similarity(video_emb, query_emb, dim=-1) * scale
+        i_sim = i_sim + torch.where(loss_msk > 0, .0, float('-inf'))
+
+        loss = 0
+
+        if 'row' in self.direction:
+            i_met = F.log_softmax(i_sim, dim=1)[batch_inds, pos_clip]
+            loss = loss - i_met.sum() / i_met.size(0)
+
+        if 'col' in self.direction:
+            j_sim = i_sim.t()
+            j_met = F.log_softmax(j_sim, dim=1)[pos_clip, batch_inds]
+            loss = loss - j_met.sum() / j_met.size(0)
+
+        loss = loss * self.loss_weight
+        return loss
+
+
+@LOSSES.register()
+class BundleLoss(nn.Module):
+
+    def __init__(self,
+                 sample_radius=1.5,
+                 loss_cls=None,
+                 loss_reg=None,
+                 loss_sal=None,
+                 loss_video_cal=None,
+                 loss_layer_cal=None):
+        super(BundleLoss, self).__init__()
+
+        self._loss_cls = build_loss(loss_cls)
+        self._loss_reg = build_loss(loss_reg)
+        self._loss_sal = build_loss(loss_sal)
+        self._loss_video_cal = build_loss(loss_video_cal)
+        self._loss_layer_cal = build_loss(loss_layer_cal)
+
+        self.sample_radius = sample_radius
+
+    def get_target_single(self, point, gt_bnd, gt_cls):
+        num_pts, num_gts = point.size(0), gt_bnd.size(0)
+
+        lens = gt_bnd[:, 1] - gt_bnd[:, 0]
+        lens = lens[None, :].repeat(num_pts, 1)
+
+        gt_seg = gt_bnd[None].expand(num_pts, num_gts, 2)
+        s = point[:, 0, None] - gt_seg[:, :, 0]
+        e = gt_seg[:, :, 1] - point[:, 0, None]
+        r_tgt = torch.stack((s, e), dim=-1)
+
+        if self.sample_radius > 0:
+            center = (gt_seg[:, :, 0] + gt_seg[:, :, 1]) / 2
+            t_mins = center - point[:, 3, None] * self.sample_radius
+            t_maxs = center + point[:, 3, None] * self.sample_radius
+            dist_s = point[:, 0, None] - torch.maximum(t_mins, gt_seg[:, :, 0])
+            dist_e = torch.minimum(t_maxs, gt_seg[:, :, 1]) - point[:, 0, None]
+            center = torch.stack((dist_s, dist_e), dim=-1)
+            cls_msk = center.min(-1)[0] >= 0
+        else:
+            cls_msk = r_tgt.min(-1)[0] >= 0
+
+        reg_dist = r_tgt.max(-1)[0]
+        reg_msk = torch.logical_and((reg_dist >= point[:, 1, None]),
+                                    (reg_dist <= point[:, 2, None]))
+
+        lens.masked_fill_(cls_msk == 0, float('inf'))
+        lens.masked_fill_(reg_msk == 0, float('inf'))
+        min_len, min_len_inds = lens.min(dim=1)
+
+        min_len_mask = torch.logical_and((lens <= (min_len[:, None] + 1e-3)),
+                                         (lens < float('inf'))).to(r_tgt.dtype)
+
+        label = F.one_hot(gt_cls[:, 0], 2).to(r_tgt.dtype)
+        c_tgt = torch.matmul(min_len_mask, label).clamp(min=0.0, max=1.0)[:, 1]
+        r_tgt = r_tgt[range(num_pts), min_len_inds] / point[:, 3, None]
+
+        return c_tgt, r_tgt
+
+    def get_target(self, data):
+        cls_tgt, reg_tgt = [], []
+
+        for i in range(data['boundary'].size(0)):
+            gt_bnd = data['boundary'][i] * data['fps'][i]
+            gt_cls = gt_bnd.new_ones(gt_bnd.size(0), 1).long()
+
+            c_tgt, r_tgt = self.get_target_single(data['point'], gt_bnd, gt_cls)
+
+            cls_tgt.append(c_tgt)
+            reg_tgt.append(r_tgt)
+
+        cls_tgt = torch.stack(cls_tgt)
+        reg_tgt = torch.stack(reg_tgt)
+
+        return cls_tgt, reg_tgt
+
+    def loss_cls(self, data, output, cls_tgt):
+        src = data['out_class'].squeeze(-1)
+        msk = torch.cat(data['pymid_msk'], dim=1)
+
+        loss_cls = self._loss_cls(src, cls_tgt, weight=msk, avg_factor=msk.sum())
+
+        output['loss_cls'] = loss_cls
+        return output
+
+    def loss_reg(self, data, output, cls_tgt, reg_tgt):
+        src = data['out_coord']
+        msk = cls_tgt.unsqueeze(2).repeat(1, 1, 2).bool()
+
+        loss_reg = self._loss_reg(src, reg_tgt, weight=msk, avg_factor=msk.sum())
+
+        output['loss_reg'] = loss_reg
+        return output
+
+    def loss_sal(self, data, output):
+        video_emb = data['video_emb']
+        query_emb = data['query_emb']
+        video_msk = data['video_msk']
+
+        saliency = data['saliency']
+        pos_clip = data['pos_clip'][:, 0]
+
+        output['loss_sal'] = self._loss_sal(video_emb, query_emb, video_msk, saliency,
+                                            pos_clip)
+        return output
+
+    def loss_cal(self, data, output):
+        pos_clip = data['pos_clip'][:, 0]
+
+        batch_inds = torch.arange(pos_clip.size(0), device=pos_clip.device)
+
+        coll_v_emb, coll_q_emb = [], []
+        for v_emb, q_emb in zip(data['coll_v'], data['coll_q']):
+            v_emb_pos = v_emb[batch_inds, pos_clip]
+            q_emb_pos = q_emb[:, 0]
+
+            coll_v_emb.append(v_emb_pos)
+            coll_q_emb.append(q_emb_pos)
+
+        v_emb = torch.stack(coll_v_emb)
+        q_emb = torch.stack(coll_q_emb)
+        output['loss_video_cal'] = self._loss_video_cal(v_emb, q_emb)
+
+        v_emb = torch.stack(coll_v_emb, dim=1)
+        q_emb = torch.stack(coll_q_emb, dim=1)
+        output['loss_layer_cal'] = self._loss_layer_cal(v_emb, q_emb)
+
+        return output
+
+    def forward(self, data, output):
+        if self._loss_reg is not None:
+            cls_tgt, reg_tgt = self.get_target(data)
+            output = self.loss_reg(data, output, cls_tgt, reg_tgt)
+        else:
+            cls_tgt = data['saliency']
+
+        if self._loss_cls is not None:
+            output = self.loss_cls(data, output, cls_tgt)
+
+        if self._loss_sal is not None:
+            output = self.loss_sal(data, output)
+
+        if self._loss_video_cal is not None or self._loss_layer_cal is not None:
+            output = self.loss_cal(data, output)
+
+        return output

models/model.py

@@ -0,0 +1,206 @@
+# Copyright (c) Ye Liu. Licensed under the BSD 3-Clause License.
+
+import math
+
+import clip
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from nncore.nn import MODELS, build_loss, build_model
+
+from .generator import PointGenerator
+
+_CLIP_ARCHS = {
+    'ViT-B/32': (768, 512, 50),
+    'ViT-B/16': (768, 512, 197),
+    'ViT-L/14': (1024, 768, 50),
+    'ViT-L/14-336px': (1024, 768, 577)
+}
+
+
+@MODELS.register()
+class R2Tuning(nn.Module):
+
+    def __init__(self,
+                 arch='ViT-B/32',
+                 init=True,
+                 dims=256,
+                 strides=(1, 2, 4, 8),
+                 buffer_size=1024,
+                 max_num_moment=50,
+                 merge_cls_sal=True,
+                 adapter_cfg=None,
+                 pyramid_cfg=None,
+                 pooling_cfg=None,
+                 class_head_cfg=None,
+                 coord_head_cfg=None,
+                 loss_cfg=None):
+        super(R2Tuning, self).__init__()
+
+        if init:
+            self.clip, _ = clip.load(arch, device='cpu')
+            for param in self.clip.parameters():
+                param.requires_grad = False
+
+        self.cfg = _CLIP_ARCHS[arch]
+        self.adapter = build_model(adapter_cfg, dims, self.cfg[:2])
+        self.pyramid = build_model(pyramid_cfg, dims, strides)
+        self.pooling = build_model(pooling_cfg, dims)
+
+        self.class_head = build_model(class_head_cfg, dims, 1)
+        self.coord_head = build_model(coord_head_cfg, dims, 2)
+
+        self.generator = PointGenerator(strides, buffer_size)
+
+        self.coef = nn.Parameter(torch.ones(len(strides)))
+        self.loss = build_loss(loss_cfg)
+
+        self.max_num_moment = max_num_moment
+        self.merge_cls_sal = merge_cls_sal
+
+    def train(self, mode=True):
+        super(R2Tuning, self).train(mode=mode)
+        if hasattr(self, 'clip'):
+            self.clip.eval()
+
+    @torch.no_grad
+    def clip_video_tower(self, video):
+        video = video.type(self.clip.dtype)
+        video = self.clip.visual.conv1(video)
+        video = video.reshape(video.size(0), video.size(1), -1).permute(0, 2, 1)
+        c_emb = video.new_zeros(video.size(0), 1, video.size(-1))
+        c_emb = self.clip.visual.class_embedding.to(video.dtype) + c_emb
+        video = torch.cat((c_emb, video), dim=1)
+        video = video + self.clip.visual.positional_embedding.to(video.dtype)
+        video = self.clip.visual.ln_pre(video).permute(1, 0, 2)
+        emb = [video]
+        for blk in self.clip.visual.transformer.resblocks:
+            emb.append(blk(emb[-1]))
+        video = torch.stack([e.permute(1, 0, 2) for e in emb])
+        return video
+
+    @torch.no_grad
+    def clip_query_tower(self, query):
+        query = self.clip.token_embedding(query).type(self.clip.dtype)
+        query = query + self.clip.positional_embedding.type(self.clip.dtype)
+        query = query.permute(1, 0, 2)
+        emb = [query]
+        for blk in self.clip.transformer.resblocks:
+            emb.append(blk(emb[-1]))
+        query = torch.stack([e.permute(1, 0, 2) for e in emb])
+        return query
+
+    def forward(self, data, mode='test'):
+        video, query = data['video'], data['query']
+
+        if hasattr(self, 'clip'):
+            video_msk = torch.where(video[:, :, 0].isfinite(), 1, 0)
+            query_msk = torch.where(query == 0, 0, 1)
+
+            video[~video.isfinite()] = 0
+
+            (b, t), d = video.size()[:2], int(math.sqrt(video.size(2) / 3))
+            video = video.view(b * t, 3, d, d)
+
+            video_emb = self.clip_video_tower(video)
+            query_emb = self.clip_query_tower(query)
+
+            n, _, p, c = video_emb.size()
+            video_emb = video_emb.view(n, b, t, p, c)
+        else:
+            video_msk = torch.where(video[:, :, 0].isfinite(), 1, 0)
+            query_msk = torch.where(query[:, :, 0].isfinite(), 1, 0)
+
+            video[~video.isfinite()] = 0
+            query[~query.isfinite()] = 0
+
+            (b, t), l = video.size()[:2], query.size(1)
+            video = video.view(b, t, -1, self.cfg[2], self.cfg[0]).permute(2, 0, 1, 3, 4)
+            query = query.view(b, l, -1, self.cfg[1]).permute(2, 0, 1, 3)
+
+            video_emb = video.float()
+            query_emb = query.float()
+
+        # video_emb: N * B * T * P * C
+        # query_emb: N * B * L * C
+
+        video_emb, query_emb, coll_v, coll_q = self.adapter(video_emb, query_emb,
+                                                            video_msk, query_msk)
+
+        pymid, pymid_msk = self.pyramid(video_emb, video_msk, return_mask=mode != 'test')
+        point = self.generator(pymid)
+
+        with torch.autocast('cuda', enabled=False):
+            video_emb = video_emb.float()
+            query_emb = self.pooling(query_emb.float(), query_msk)
+
+            out_class = [self.class_head(e.float()) for e in pymid]
+            out_class = torch.cat(out_class, dim=1)
+
+            if self.coord_head is not None:
+                out_coord = [
+                    self.coord_head(e.float()).exp() * self.coef[i]
+                    for i, e in enumerate(pymid)
+                ]
+                out_coord = torch.cat(out_coord, dim=1)
+            else:
+                out_coord = None
+
+            output = dict(_avg_factor=b)
+
+            if mode != 'test':
+                data['coll_v'] = [e.float() for e in coll_v]
+                data['coll_q'] = [self.pooling(e.float(), query_msk) for e in coll_q]
+
+                data['point'] = point
+                data['video_emb'] = video_emb
+                data['query_emb'] = query_emb
+                data['video_msk'] = video_msk
+                data['pymid_msk'] = pymid_msk
+                data['out_class'] = out_class
+                data['out_coord'] = out_coord
+
+                output = self.loss(data, output)
+
+            if mode != 'train':
+                assert b == 1, 'batch size larger than 1 is not supported for inference'
+                out_class = out_class.sigmoid()
+                out_score = F.cosine_similarity(video_emb, query_emb, dim=-1)
+
+                output['_out'] = dict(label=data.get('label', [None])[0])
+
+                pyd_shape = [e.size(1) for e in pymid]
+                pyd_class = out_class[0, :, 0].split(pyd_shape)
+
+                saliency = []
+                for shape, score in zip(pyd_shape, pyd_class):
+                    if t >= shape:
+                        score = score.repeat_interleave(int(t / shape))
+                        postfix = score[-1:].repeat(t - score.size(0))
+                        score = torch.cat((score, postfix))
+                    else:
+                        scale = int(shape / t)
+                        score = F.max_pool1d(score.unsqueeze(0), scale, stride=scale)[0]
+                    saliency.append(score)
+
+                saliency = torch.stack(saliency).amax(dim=0)
+
+                if self.merge_cls_sal:
+                    saliency *= out_score[0]
+
+                output['_out']['saliency'] = saliency
+
+                if self.coord_head is not None:
+                    boundary = out_coord[0]
+                    boundary[:, 0] *= -1
+                    boundary *= point[:, 3, None].repeat(1, 2)
+                    boundary += point[:, 0, None].repeat(1, 2)
+                    boundary /= data['fps'][0]
+                    boundary = torch.cat((boundary, out_class[0]), dim=-1)
+
+                    _, inds = out_class[0, :, 0].sort(descending=True)
+                    boundary = boundary[inds[:self.max_num_moment]]
+
+                    output['_out']['boundary'] = boundary
+
+        return output

requirements.txt

@@ -0,0 +1,6 @@
+git+https://github.com/openai/CLIP.git@a1d0717
+decord==0.6.0
+matplotlib==3.9.0
+nncore==0.4.3
+torch==2.2.1
+torchvision==0.17.1

setup.cfg

@@ -0,0 +1,15 @@
+[yapf]
+column_limit = 90
+based_on_style = pep8
+blank_line_before_nested_class_or_def = true
+split_before_expression_after_opening_paren = true
+
+[isort]
+line_length = 90
+multi_line_output = 0
+known_third_party = clip,decord,gradio,nncore,numpy,torch,torchvision
+no_lines_before = STDLIB,LOCALFOLDER
+default_section = FIRSTPARTY
+
+[flake8]
+max-line-length = 90