Spacetimegpt

🚀 SpaceTimeGPT - 视频字幕生成模型

SpaceTimeGPT 是一个能够进行空间和时间推理的视频描述生成模型。它可以对输入的视频进行分析，并输出描述视频中事件的句子。

🔍 项目信息

📊 模型评估结果

🚀 快速开始

(部分图表来自 1, 2, 3)

SpaceTimeGPT 是一个能够进行空间和时间推理的视频描述生成模型。给定一个视频，模型会采样并分析八帧图像，然后通过自回归生成视频中所发生事件的句子描述。

🏗️ 架构与训练

• 视觉编码器：timesformer-base-finetuned-k600
• 文本解码器：gpt2

编码器和解码器分别使用预训练的视频分类和句子补全权重进行初始化。通过编码器 - 解码器交叉注意力机制来统一视觉和语言领域。该模型在视频字幕生成任务上进行了端到端的微调。更多详细信息请参考 GitHub 仓库。

💻 使用示例

基础用法

import av
import numpy as np
import torch
from transformers import AutoImageProcessor, AutoTokenizer, VisionEncoderDecoderModel

device = "cuda" if torch.cuda.is_available() else "cpu"

# 加载预训练的处理器、分词器和模型
image_processor = AutoImageProcessor.from_pretrained("MCG-NJU/videomae-base")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = VisionEncoderDecoderModel.from_pretrained("Neleac/timesformer-gpt2-video-captioning").to(device)

# 加载视频
video_path = "never_gonna_give_you_up.mp4"
container = av.open(video_path)

# 从视频中提取均匀间隔的帧
seg_len = container.streams.video[0].frames
clip_len = model.config.encoder.num_frames
indices = set(np.linspace(0, seg_len, num=clip_len, endpoint=False).astype(np.int64))
frames = []
container.seek(0)
for i, frame in enumerate(container.decode(video=0)):
    if i in indices:
        frames.append(frame.to_ndarray(format="rgb24"))

# 生成字幕
gen_kwargs = {
    "min_length": 10, 
    "max_length": 20, 
    "num_beams": 8,
}
pixel_values = image_processor(frames, return_tensors="pt").pixel_values.to(device)
tokens = model.generate(pixel_values, **gen_kwargs)
caption = tokenizer.batch_decode(tokens, skip_special_tokens=True)[0]
print(caption) # 一名男子和一名女子在镜子前的舞台上跳舞。

🧑‍💻 作者信息

• 👾 Discord
• 🐙 GitHub
• 🤝 LinkedIn