|
|
|
|
|
import random |
|
from functools import partial |
|
|
|
import clip |
|
import decord |
|
import gradio as gr |
|
import nncore |
|
import numpy as np |
|
import torch |
|
import torchvision.transforms.functional as F |
|
from decord import VideoReader |
|
from nncore.engine import load_checkpoint |
|
from nncore.nn import build_model |
|
|
|
import pandas as pd |
|
|
|
TITLE = '🌀R2-Tuning: Efficient Image-to-Video Transfer Learning for Video Temporal Grounding' |
|
|
|
TITLE_MD = '<h1 align="center">🌀R<sup>2</sup>-Tuning: Efficient Image-to-Video Transfer Learning for Video Temporal Grounding</h1>' |
|
DESCRIPTION_MD = 'R<sup>2</sup>-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>.' |
|
GUIDE_MD = '### User Guide:\n1. Upload a video or click "random" to sample one.\n2. Input a text query. A good practice is to write a sentence with 5~15 words.\n3. Click "submit" and you\'ll see the moment retrieval and highlight detection results on the right.' |
|
|
|
CONFIG = 'configs/qvhighlights/r2_tuning_qvhighlights.py' |
|
WEIGHT = 'https://huggingface.co/yeliudev/R2-Tuning/resolve/main/checkpoints/r2_tuning_qvhighlights-ed516355.pth' |
|
|
|
|
|
EXAMPLES = [ |
|
('data/gTAvxnQtjXM_60.0_210.0.mp4', 'A man in a white t shirt wearing a backpack is showing a nearby cathedral.'), |
|
('data/pA6Z-qYhSNg_210.0_360.0.mp4', 'Different Facebook posts on transgender bathrooms are shown.'), |
|
('data/CkWOpyrAXdw_210.0_360.0.mp4', 'Indian girl cleaning her kitchen before cooking.'), |
|
('data/ocLUzCNodj4_360.0_510.0.mp4', 'A woman stands in her bedroom in front of a mirror and talks.'), |
|
('data/HkLfNhgP0TM_660.0_810.0.mp4', 'Woman lays down on the couch while talking to the camera.') |
|
] |
|
|
|
|
|
|
|
def convert_time(seconds): |
|
minutes, seconds = divmod(round(max(seconds, 0)), 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()) * 0.9 + 0.05).tolist() |
|
hd = pd.DataFrame(dict(x=range(0, len(hd) * 2, 2), y=hd)) |
|
|
|
return mr, hd |
|
|
|
|
|
model, cfg = init_model(CONFIG, WEIGHT) |
|
|
|
fn = partial(main, model=model, cfg=cfg) |
|
|
|
with gr.Blocks(title=TITLE) as demo: |
|
gr.Markdown(TITLE_MD) |
|
gr.Markdown(DESCRIPTION_MD) |
|
gr.Markdown(GUIDE_MD) |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
video = gr.Video(label='Video') |
|
query = gr.Textbox(label='Text Query') |
|
|
|
with gr.Row(): |
|
random_btn = gr.Button(value='🔮 Random') |
|
gr.ClearButton([video, query], value='🗑️ Reset') |
|
submit_btn = gr.Button(value='🚀 Submit') |
|
|
|
with gr.Column(): |
|
mr = gr.DataFrame( |
|
headers=['Start Time', 'End Time', 'Score'], label='Moment Retrieval') |
|
hd = gr.LinePlot( |
|
x='x', |
|
y='y', |
|
x_title='Time (seconds)', |
|
y_title='Saliency Score', |
|
label='Highlight Detection') |
|
|
|
random_btn.click(lambda: random.sample(EXAMPLES, 1)[0], None, [video, query]) |
|
submit_btn.click(fn, [video, query], [mr, hd]) |
|
|
|
demo.launch() |
|
|
