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