SpaceTimeGPT is a video description generation model capable of spatial and temporal reasoning. Given a video, eight frames are sampled and analyzed by the model. The output is a sentence description of the events that occured in the video, generated using autoregression.
The encoder and decoder are initialized using pretrained weights for video classification and sentence completion, respectively. Encoder-decoder cross attention is used to unify the visual and linguistic domains. The model is fine-tuned end-to-end on the video captioning task. See
GitHub repository for details.
1import av
2import numpy as np
3import torch
4from transformers import AutoImageProcessor, AutoTokenizer, VisionEncoderDecoderModel
5
6device = "cuda" if torch.cuda.is_available() else "cpu"
7
8# load pretrained processor, tokenizer, and model
9image_processor = AutoImageProcessor.from_pretrained("MCG-NJU/videomae-base")
10tokenizer = AutoTokenizer.from_pretrained("gpt2")
11model = VisionEncoderDecoderModel.from_pretrained("Neleac/timesformer-gpt2-video-captioning").to(device)
12
13# load video
14video_path = "never_gonna_give_you_up.mp4"
15container = av.open(video_path)
16
17# extract evenly spaced frames from video
18seg_len = container.streams.video[0].frames
19clip_len = model.config.encoder.num_frames
20indices = set(np.linspace(0, seg_len, num=clip_len, endpoint=False).astype(np.int64))
21frames = []
22container.seek(0)
23for i, frame in enumerate(container.decode(video=0)):
24 if i in indices:
25 frames.append(frame.to_ndarray(format="rgb24"))
26
27# generate caption
28gen_kwargs = {
29 "min_length": 10,
30 "max_length": 20,
31 "num_beams": 8,
32}
33pixel_values = image_processor(frames, return_tensors="pt").pixel_values.to(device)
34tokens = model.generate(pixel_values, **gen_kwargs)
35caption = tokenizer.batch_decode(tokens, skip_special_tokens=True)[0]
36print(caption) # A man and a woman are dancing on a stage in front of a mirror.
