233 lines
7.8 KiB
Python
233 lines
7.8 KiB
Python
"""
|
|
Pluggable Feature Extractors for FVD Computation.
|
|
Supports I3D (standard), CLIP, and VideoMAE via a common interface.
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from abc import ABC, abstractmethod
|
|
from huggingface_hub import hf_hub_download
|
|
from tqdm import tqdm
|
|
|
|
try:
|
|
from transformers import CLIPModel, CLIPProcessor, VideoMAEModel
|
|
TRANSFORMERS_AVAILABLE = True
|
|
except ImportError:
|
|
TRANSFORMERS_AVAILABLE = False
|
|
|
|
|
|
class BaseFeatureExtractor(ABC, nn.Module):
|
|
"""Abstract base class for all video feature extractors."""
|
|
|
|
def __init__(self, device: str = 'cuda'):
|
|
super().__init__()
|
|
self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
|
|
|
|
@property
|
|
@abstractmethod
|
|
def feature_dim(self) -> int:
|
|
"""Dimension of the output feature vector."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def preprocess(self, videos: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
videos: [B, T, C, H, W] in [0, 255] range.
|
|
Returns:
|
|
Preprocessed tensor ready for the model.
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def extract_features_batch(self, videos: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Extract features for a single batch.
|
|
Args:
|
|
videos: [B, T, C, H, W] (raw input)
|
|
Returns:
|
|
Features: [B, feature_dim]
|
|
"""
|
|
pass
|
|
|
|
@torch.no_grad()
|
|
def extract_features(self, videos: torch.Tensor, batch_size: int = 32, verbose: bool = True) -> torch.Tensor:
|
|
"""
|
|
Extract features for a large tensor of videos by batching.
|
|
"""
|
|
N = len(videos)
|
|
all_features = []
|
|
|
|
iterator = range(0, N, batch_size)
|
|
if verbose:
|
|
iterator = tqdm(iterator, desc=f"Extracting features ({self.__class__.__name__})")
|
|
|
|
for i in iterator:
|
|
batch = videos[i:i + batch_size].to(self.device)
|
|
features = self.extract_features_batch(batch)
|
|
all_features.append(features.cpu())
|
|
|
|
return torch.cat(all_features, dim=0)
|
|
|
|
|
|
# 1. I3D Extractor (The Standard FVD Metric)
|
|
class I3DFeatureExtractor(BaseFeatureExtractor):
|
|
REPO_ID = 'flateon/FVD-I3D-torchscript'
|
|
MODEL_FILENAME = 'i3d_torchscript.pt'
|
|
|
|
def __init__(self, device: str = 'cuda', cache_dir: str | None = None):
|
|
super().__init__(device)
|
|
self.cache_dir = cache_dir
|
|
self.model = self._load_model()
|
|
self.model.eval()
|
|
self.model.to(self.device)
|
|
|
|
@property
|
|
def feature_dim(self) -> int:
|
|
return 400
|
|
|
|
def _load_model(self) -> torch.nn.Module:
|
|
try:
|
|
model_path = hf_hub_download(repo_id=self.REPO_ID, filename=self.MODEL_FILENAME, cache_dir=self.cache_dir)
|
|
return torch.jit.load(model_path, map_location=self.device)
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to load I3D model: {e}") from e
|
|
|
|
def preprocess(self, videos: torch.Tensor) -> torch.Tensor:
|
|
"""Standard I3D preprocessing: Resize to 224, Norm to [-1, 1]."""
|
|
B, T, C, H, W = videos.shape
|
|
|
|
if T < 10:
|
|
raise ValueError(f"I3D requires at least 10 frames, got {T}")
|
|
|
|
# Normalize to [0, 1]
|
|
if videos.max() > 1.0:
|
|
videos = videos / 255.0
|
|
|
|
# Scale to [-1, 1]
|
|
videos = videos * 2.0 - 1.0
|
|
|
|
# Resize to 224x224
|
|
if H != 224 or W != 224:
|
|
videos = videos.reshape(B * T, C, H, W)
|
|
videos = F.interpolate(videos, size=(224, 224), mode='bilinear', align_corners=False)
|
|
videos = videos.reshape(B, T, C, 224, 224)
|
|
|
|
# [B, T, C, H, W] -> [B, C, T, H, W]
|
|
return videos.permute(0, 2, 1, 3, 4).contiguous()
|
|
|
|
def extract_features_batch(self, videos: torch.Tensor) -> torch.Tensor:
|
|
batch = self.preprocess(videos)
|
|
# TorchScript I3D returns raw logits when return_features=True
|
|
return self.model(batch, rescale=False, resize=False, return_features=True)
|
|
|
|
|
|
# 2. CLIP Extractor (Semantic/Content Quality)
|
|
class CLIPFeatureExtractor(BaseFeatureExtractor):
|
|
|
|
def __init__(self, device: str = 'cuda', model_name: str = "openai/clip-vit-base-patch32"):
|
|
if not TRANSFORMERS_AVAILABLE:
|
|
raise ImportError("Please install transformers: uv pip install transformers")
|
|
super().__init__(device)
|
|
self.processor = CLIPProcessor.from_pretrained(model_name)
|
|
self.model = CLIPModel.from_pretrained(model_name).to(self.device)
|
|
self.model.eval()
|
|
self._feature_dim = self.model.config.projection_dim
|
|
|
|
@property
|
|
def feature_dim(self) -> int:
|
|
return self._feature_dim
|
|
|
|
def preprocess(self, videos: torch.Tensor) -> torch.Tensor:
|
|
# Ensure values are [0, 255]
|
|
if videos.max() <= 1.0:
|
|
videos = videos * 255.0
|
|
|
|
return videos.to(torch.uint8)
|
|
|
|
def extract_features_batch(self, videos: torch.Tensor) -> torch.Tensor:
|
|
# Input: [B, T, C, H, W]
|
|
B, T, C, H, W = videos.shape
|
|
videos = self.preprocess(videos)
|
|
|
|
# Flatten B*T to treat frames as images
|
|
images = videos.view(B * T, C, H, W)
|
|
|
|
# HF Processor
|
|
inputs = self.processor(images=images, return_tensors="pt", padding=True)
|
|
inputs = {k: v.to(self.device) for k, v in inputs.items()}
|
|
|
|
# Extract features [B*T, Dim]
|
|
outputs = self.model.get_image_features(**inputs)
|
|
|
|
# Reshape [B, T, Dim] and Average Pooling over time
|
|
outputs = outputs.view(B, T, -1)
|
|
return outputs.mean(dim=1)
|
|
|
|
|
|
# 3. VideoMAE Extractor (Structure/Motion Quality)
|
|
class VideoMAEFeatureExtractor(BaseFeatureExtractor):
|
|
|
|
def __init__(self, device: str = 'cuda', model_name: str = "MCG-NJU/videomae-base"):
|
|
if not TRANSFORMERS_AVAILABLE:
|
|
raise ImportError("Please install transformers: uv pip install transformers")
|
|
super().__init__(device)
|
|
self.model = VideoMAEModel.from_pretrained(model_name).to(self.device)
|
|
self.model.eval()
|
|
|
|
self.register_buffer('mean', torch.tensor([0.485, 0.456, 0.406], device=self.device).view(1, 1, 3, 1, 1))
|
|
self.register_buffer('std', torch.tensor([0.229, 0.224, 0.225], device=self.device).view(1, 1, 3, 1, 1))
|
|
|
|
@property
|
|
def feature_dim(self) -> int:
|
|
return self.model.config.hidden_size
|
|
|
|
def preprocess(self, videos: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Efficient GPU-based preprocessing.
|
|
Input: [B, T, C, H, W] in range [0, 255]
|
|
"""
|
|
B, T, C, H, W = videos.shape
|
|
|
|
# 1. Resize to 224x224
|
|
if H != 224 or W != 224:
|
|
videos = videos.view(B * T, C, H, W)
|
|
videos = F.interpolate(videos, size=(224, 224), mode='bilinear', align_corners=False)
|
|
videos = videos.view(B, T, C, 224, 224)
|
|
|
|
# 2. Normalize to [0, 1]
|
|
if videos.dtype != torch.float32:
|
|
videos = videos.float()
|
|
|
|
if videos.max() > 1.0:
|
|
videos = videos / 255.0
|
|
|
|
# 3. Apply ImageNet Mean/Std
|
|
return (videos - self.mean) / self.std
|
|
|
|
def extract_features_batch(self, videos: torch.Tensor) -> torch.Tensor:
|
|
# Input: [B, T, C, H, W]
|
|
|
|
# Fast GPU Preprocessing
|
|
pixel_values = self.preprocess(videos)
|
|
|
|
# Forward pass
|
|
outputs = self.model(pixel_values)
|
|
|
|
# Global Average Pooling of last hidden state [B, T_patches, 768] -> [B, 768]
|
|
return outputs.last_hidden_state.mean(dim=1)
|
|
|
|
|
|
# Factory
|
|
def load_extractor(name: str, device: str = 'cuda') -> BaseFeatureExtractor:
|
|
name = name.lower()
|
|
if name == 'i3d':
|
|
return I3DFeatureExtractor(device)
|
|
elif name == 'clip':
|
|
return CLIPFeatureExtractor(device)
|
|
elif name == 'videomae':
|
|
return VideoMAEFeatureExtractor(device)
|
|
else:
|
|
raise ValueError(f"Unknown extractor: {name}. Options: i3d, clip, videomae")
|