Video Understanding

import torch
import torchvision.io as tvio
import torchvision.transforms.functional as TF
import torchvision.models.optical_flow as of_models
from torchvision.models.optical_flow import Raft_Large_Weights
from PIL import Image
import numpy as np

def load_video_frames(path: str, n_frames=16, target_size=(224, 224)):
    """Load uniformly sampled frames from a video file."""
    video, _, _ = tvio.read_video(path, pts_unit='sec')  # (T, H, W, C) uint8
    T = video.shape[0]
    indices = torch.linspace(0, T-1, n_frames).long()
    frames = video[indices].permute(0, 3, 1, 2).float() / 255.0  # (T, C, H, W)
    frames = torch.stack([TF.resize(f, list(target_size)) for f in frames])
    mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
    std  = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
    return (frames - mean) / std

def augment_clip(frames: torch.Tensor) -> torch.Tensor:
    """Consistent spatial augmentation across all frames."""
    T, C, H, W = frames.shape
    i = torch.randint(0, H-224, (1,)).item()
    j = torch.randint(0, W-224, (1,)).item()
    frames = frames[:, :, i:i+224, j:j+224]
    if torch.rand(1) > 0.5: frames = frames.flip(-1)
    return frames

# Optical flow with RAFT
raft = of_models.raft_large(weights=Raft_Large_Weights.C_T_SKHT_V2).eval()

def compute_flow(frame1: Image.Image, frame2: Image.Image) -> np.ndarray:
    to_tensor = lambda img: TF.to_tensor(img.convert("RGB")).unsqueeze(0)
    with torch.no_grad():
        flows = raft(to_tensor(frame1), to_tensor(frame2))
    return flows[-1].squeeze(0).permute(1, 2, 0).numpy()  # (H, W, 2)

import torch
from transformers import (
    VideoMAEImageProcessor, VideoMAEForVideoClassification,
    TrainingArguments, Trainer,
)
from torch.utils.data import Dataset
import torchvision.io as tvio

# Load pre-trained VideoMAE (Kinetics-400)
processor = VideoMAEImageProcessor.from_pretrained("MCG-NJU/videomae-base-finetuned-kinetics")
model = VideoMAEForVideoClassification.from_pretrained(
    "MCG-NJU/videomae-base",   # base without head
    num_labels=10,              # your class count
    ignore_mismatched_sizes=True,
)

class VideoDataset(Dataset):
    def __init__(self, video_paths, labels, processor, n_frames=16):
        self.video_paths = video_paths
        self.labels = labels
        self.processor = processor
        self.n_frames = n_frames

    def __len__(self): return len(self.video_paths)

    def __getitem__(self, idx):
        video, _, _ = tvio.read_video(self.video_paths[idx], pts_unit='sec')
        T = video.shape[0]
        indices = torch.linspace(0, T-1, self.n_frames).long()
        frames = [video[i].numpy() for i in indices]   # list of HxWxC uint8
        encoding = self.processor(frames, return_tensors="pt")
        return {
            'pixel_values': encoding['pixel_values'].squeeze(0),
            'labels': torch.tensor(self.labels[idx]),
        }

args = TrainingArguments(
    output_dir="videomae-custom",
    per_device_train_batch_size=4, num_train_epochs=20,
    learning_rate=5e-5, warmup_ratio=0.1, weight_decay=0.05,
    evaluation_strategy="epoch", fp16=True,
)

# Multi-clip inference for long videos
@torch.no_grad()
def predict_long_video(model, processor, video_path, n_clips=10, n_frames=16, device='cuda'):
    video, _, _ = tvio.read_video(video_path, pts_unit='sec')
    T, clip_len = video.shape[0], video.shape[0] // n_clips
    model.eval().to(device)
    all_probs = []
    for i in range(n_clips):
        clip = video[i*clip_len:(i+1)*clip_len]
        idxs = torch.linspace(0, len(clip)-1, n_frames).long()
        frames = [clip[j].numpy() for j in idxs]
        inputs = processor(frames, return_tensors="pt")
        logits = model(pixel_values=inputs['pixel_values'].to(device)).logits[0]
        all_probs.append(logits.softmax(dim=-1).cpu())
    avg = torch.stack(all_probs).mean(0)
    pred = model.config.id2label[avg.argmax().item()]
    print(f"Predicted: {pred} ({avg.max():.3f})")
    return pred, avg