from decord import VideoReader, cpu  
import torch  
import numpy as np  

from transformers import VideoMAEFeatureExtractor, VideoMAEForVideoClassification  
from huggingface_hub import hf_hub_download  

np.random.seed(0)  

def sample_frame_indices(clip_len, frame_sample_rate, seg_len):  
    converted_len = int(clip_len * frame_sample_rate)  
    end_idx = np.random.randint(converted_len, seg_len)  
    start_idx = end_idx - converted_len  
    indices = np.linspace(start_idx, end_idx, num=clip_len)  
    indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64)  
    return indices  

# 视频片段包含300帧（30 FPS下10秒）  
file_path = hf_hub_download(  
    repo_id="nateraw/dino-clips", filename="archery.mp4", repo_type="space"  
)  
videoreader = VideoReader(file_path, num_threads=1, ctx=cpu(0))  

# 采样16帧  
videoreader.seek(0)  
indices = sample_frame_indices(clip_len=16, frame_sample_rate=4, seg_len=len(videoreader))  
video = videoreader.get_batch(indices).asnumpy()  

feature_extractor = VideoMAEFeatureExtractor.from_pretrained("nateraw/videomae-base-finetuned-ucf101")  
model = VideoMAEForVideoClassification.from_pretrained("nateraw/videomae-base-finetuned-ucf101")  

inputs = feature_extractor(list(video), return_tensors="pt")  

with torch.no_grad():  
    outputs = model(**inputs)  
    logits = outputs.logits  

# 模型预测101个UCF101类别之一  
predicted_label = logits.argmax(-1).item()  
print(model.config.id2label[predicted_label])