385 lines
14 KiB
Python
385 lines
14 KiB
Python
import numpy as np
|
|
import scipy.linalg
|
|
import torch
|
|
from pathlib import Path
|
|
from collections.abc import Iterator
|
|
import pickle
|
|
from dataclasses import dataclass, field
|
|
from .feature_extractors import BaseFeatureExtractor, load_extractor
|
|
from .video_utils import ClipSamplingStrategy, load_video_clips_streaming
|
|
|
|
|
|
def compute_statistics(features: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Compute mean and covariance."""
|
|
mu = np.mean(features, axis=0)
|
|
sigma = np.cov(features, rowvar=False)
|
|
return mu, sigma
|
|
|
|
|
|
def compute_frechet_distance(mu1: np.ndarray,
|
|
sigma1: np.ndarray,
|
|
mu2: np.ndarray,
|
|
sigma2: np.ndarray,
|
|
eps: float = 1e-6) -> float:
|
|
"""
|
|
Compute Fréchet distance between two Gaussians.
|
|
"""
|
|
sigma1 = sigma1 + eps * np.eye(sigma1.shape[0])
|
|
sigma2 = sigma2 + eps * np.eye(sigma2.shape[0])
|
|
|
|
diff = mu1 - mu2
|
|
mean_distance = np.sum(diff**2)
|
|
|
|
trace_sum = np.trace(sigma1 + sigma2)
|
|
|
|
covmean = scipy.linalg.sqrtm(sigma1 @ sigma2)
|
|
|
|
if np.iscomplexobj(covmean):
|
|
if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
|
|
print(f"Warning: Imaginary component: {np.max(np.abs(covmean.imag))}")
|
|
covmean = covmean.real
|
|
|
|
trace_product = np.trace(covmean)
|
|
|
|
fvd = mean_distance + trace_sum - 2 * trace_product
|
|
|
|
return float(fvd)
|
|
|
|
|
|
@dataclass
|
|
class FVDConfig:
|
|
# default configuration for FVD computation:
|
|
|
|
# Video selection
|
|
num_videos: int = 2048
|
|
|
|
# Feature Extractor Selection
|
|
extractor_model: str = 'i3d' # Options: 'i3d', 'clip', 'videomae'
|
|
|
|
# Clip sampling
|
|
num_frames_per_clip: int = 16
|
|
num_clips_per_video: int = 1
|
|
clip_strategy: str | ClipSamplingStrategy = 'beginning'
|
|
|
|
# Temporal subsampling
|
|
frame_stride: int = 1 # 1=no subsampling, 2=every 2nd, 8=every 8th
|
|
temporal_stride: int = 1 # For sliding window clips
|
|
|
|
# Data processing
|
|
video_extensions: list[str] = field(default_factory=lambda: ['.mp4', '.avi', '.mov', '.mkv'])
|
|
support_frame_dirs: bool = True
|
|
|
|
# Computation
|
|
batch_size: int = 32
|
|
device: str = 'cuda'
|
|
|
|
use_streaming: bool = True
|
|
resize_before_extraction: bool = True
|
|
|
|
# Caching
|
|
cache_real_features: str | None = None
|
|
i3d_model_path: str | None = None
|
|
|
|
# Reproducibility
|
|
seed: int | None = None
|
|
|
|
@classmethod
|
|
def fvd2048_16f(cls) -> 'FVDConfig':
|
|
"""Standard FVD protocol: 2048 videos, 16 frames, beginning clip."""
|
|
return cls(num_videos=2048, num_frames_per_clip=16, clip_strategy='beginning', use_streaming=True)
|
|
|
|
@classmethod
|
|
def fvd2048_128f(cls) -> 'FVDConfig':
|
|
"""Long video protocol: 2048 videos, 128 frames."""
|
|
return cls(num_videos=2048, num_frames_per_clip=128, clip_strategy='beginning', use_streaming=True)
|
|
|
|
@classmethod
|
|
def quick_test(cls) -> 'FVDConfig':
|
|
"""Quick test config: 100 videos, 16 frames."""
|
|
return cls(num_videos=100, num_frames_per_clip=16, clip_strategy='beginning')
|
|
|
|
def to_dict(self) -> dict:
|
|
"""Export config to dict for logging"""
|
|
d = self.__dict__.copy()
|
|
d['clip_strategy'] = str(self.clip_strategy)
|
|
return d
|
|
|
|
def __str__(self) -> str:
|
|
"""Human-readable protocol name"""
|
|
desc = f"FVD_{self.extractor_model.upper()}_{self.num_videos}_{self.num_frames_per_clip}f"
|
|
if self.frame_stride > 1:
|
|
desc += f"_subsample{self.frame_stride}"
|
|
if self.num_clips_per_video > 1:
|
|
desc += f"_{self.num_clips_per_video}clips"
|
|
if self.clip_strategy != 'beginning':
|
|
desc += f"_{self.clip_strategy}"
|
|
return desc
|
|
|
|
|
|
def extract_features_streaming(video_generator: Iterator[torch.Tensor],
|
|
extractor: BaseFeatureExtractor,
|
|
batch_size: int = 32,
|
|
max_clips: int | None = None,
|
|
verbose: bool = True) -> np.ndarray:
|
|
"""
|
|
Extract features from a video clip generator using streaming.
|
|
"""
|
|
all_features = []
|
|
batch = []
|
|
|
|
if verbose:
|
|
print(f"Extracting features with batch_size={batch_size}...")
|
|
|
|
with torch.no_grad():
|
|
for clip_count, clip in enumerate(video_generator):
|
|
batch.append(clip)
|
|
|
|
# Process batch when full
|
|
if len(batch) == batch_size:
|
|
batch_tensor = torch.stack(batch).to(extractor.device)
|
|
features = extractor.extract_features_batch(batch_tensor)
|
|
|
|
all_features.append(features.detach().cpu().numpy())
|
|
batch = []
|
|
|
|
if verbose and clip_count % (batch_size * 10) == 0:
|
|
print(f"Processed {clip_count} clips...")
|
|
|
|
if max_clips is not None and clip_count >= max_clips:
|
|
break
|
|
|
|
# Process remaining clips
|
|
if len(batch) > 0:
|
|
batch_tensor = torch.stack(batch).to(extractor.device)
|
|
features = extractor.extract_features_batch(batch_tensor)
|
|
all_features.append(features.detach().cpu().numpy())
|
|
|
|
if len(all_features) == 0:
|
|
raise RuntimeError("No features extracted - check video loading")
|
|
|
|
features = np.concatenate(all_features, axis=0)
|
|
|
|
if verbose:
|
|
print(f"Extracted {len(features)} feature vectors")
|
|
|
|
return features
|
|
|
|
|
|
def load_or_compute_features(videos: str | Path | torch.Tensor,
|
|
extractor: BaseFeatureExtractor,
|
|
config: FVDConfig,
|
|
cache_path: str | None = None,
|
|
cache_name: str = "real_features") -> np.ndarray:
|
|
"""Load features from cache or compute (with streaming support)"""
|
|
|
|
if cache_path is not None:
|
|
script_dir = Path(__file__).parent
|
|
cache_dir = script_dir / cache_path
|
|
cache_file = cache_dir / f"{config.extractor_model}_{cache_name}.pkl"
|
|
|
|
if cache_file.exists():
|
|
print(f"Loading cached features from {cache_file}")
|
|
with open(cache_file, 'rb') as f:
|
|
features = pickle.load(f)
|
|
|
|
# Validate and limit based on config
|
|
max_features = config.num_videos * config.num_clips_per_video
|
|
|
|
if len(features) < max_features:
|
|
print(f"WARNING: Cache has {len(features)} features but need {max_features}")
|
|
print("Cached features insufficient - will recompute...")
|
|
elif len(features) > max_features:
|
|
print(f"Using {max_features} features from cache (truncated from {len(features)})")
|
|
features = features[:max_features]
|
|
return features
|
|
else:
|
|
print(f"Using all {len(features)} cached features")
|
|
return features
|
|
|
|
print("Computing features from scratch...")
|
|
|
|
if isinstance(videos, (str | Path)):
|
|
target_size = (224, 224) if config.resize_before_extraction else None
|
|
|
|
video_generator = load_video_clips_streaming(videos,
|
|
num_frames=config.num_frames_per_clip,
|
|
max_videos=config.num_videos,
|
|
clip_strategy=config.clip_strategy,
|
|
frame_stride=config.frame_stride,
|
|
num_clips_per_video=config.num_clips_per_video,
|
|
video_extensions=config.video_extensions,
|
|
support_frame_dirs=config.support_frame_dirs,
|
|
target_size=target_size,
|
|
verbose=True)
|
|
|
|
max_clips = config.num_videos * config.num_clips_per_video
|
|
features = extract_features_streaming(video_generator,
|
|
extractor,
|
|
batch_size=config.batch_size,
|
|
max_clips=max_clips,
|
|
verbose=True)
|
|
else:
|
|
print(f"Extracting features from {len(videos)} video tensors...")
|
|
features = extractor.extract_features(videos, batch_size=config.batch_size, verbose=True)
|
|
features = features.numpy()
|
|
|
|
# Validate feature count
|
|
expected_count = config.num_videos * config.num_clips_per_video
|
|
if len(features) < expected_count:
|
|
raise ValueError(f"ERROR: Only extracted {len(features)} features, but need {expected_count}!\n"
|
|
f"Found fewer videos than expected. Check your video directory.")
|
|
elif len(features) > expected_count:
|
|
print(f"Truncating {len(features)} features to {expected_count}")
|
|
features = features[:expected_count]
|
|
|
|
# Cache features if requested
|
|
if cache_path is not None:
|
|
script_dir = Path(__file__).parent
|
|
cache_dir = script_dir / cache_path
|
|
cache_dir.mkdir(parents=True, exist_ok=True)
|
|
cache_file = cache_dir / f"{config.extractor_model}_{cache_name}.pkl"
|
|
print(f"Caching features to {cache_file}")
|
|
with open(cache_file, 'wb') as f:
|
|
pickle.dump(features, f)
|
|
|
|
return features
|
|
|
|
|
|
def compute_fvd_with_config(real_videos: str | Path | torch.Tensor,
|
|
gen_videos: str | Path | torch.Tensor,
|
|
config: FVDConfig,
|
|
verbose: bool = True) -> dict:
|
|
"""
|
|
Compute FVD using a standardized configuration.
|
|
|
|
This is the recommended way to compute FVD for reproducibility.
|
|
|
|
Args:
|
|
real_videos: Path or tensors
|
|
gen_videos: Path or tensors
|
|
config: FVDConfig specifying protocol
|
|
verbose: Print progress
|
|
|
|
Returns:
|
|
results: Dictionary with:
|
|
- 'fvd': FVD score (float)
|
|
- 'protocol': Protocol name (str)
|
|
- 'model': Feature extractor model name (str)
|
|
- 'config': Configuration dict
|
|
|
|
Example:
|
|
>>> config = FVDConfig.fvd2048_16f()
|
|
>>> results = compute_fvd_with_config('data/real/', 'outputs/gen/', config)
|
|
>>> print(f"FVD: {results['fvd']:.2f}")
|
|
"""
|
|
|
|
# Seed for reproducibility
|
|
if config.seed is not None:
|
|
import random as _rnd
|
|
_rnd.seed(config.seed)
|
|
np.random.seed(config.seed)
|
|
torch.manual_seed(config.seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed_all(config.seed)
|
|
|
|
if verbose:
|
|
print("=" * 70)
|
|
print(f"Computing FVD with protocol: {config}")
|
|
print(f"Model: {config.extractor_model.upper()}")
|
|
print("=" * 70)
|
|
print("\nConfiguration:")
|
|
for key, value in config.to_dict().items():
|
|
print(f" {key}: {value}")
|
|
print()
|
|
|
|
# Initialize Extractor using Factory
|
|
if verbose:
|
|
print(f"\nInitializing {config.extractor_model.upper()} model on {config.device}...")
|
|
|
|
extractor = load_extractor(config.extractor_model, device=config.device)
|
|
|
|
# Extract features
|
|
if verbose:
|
|
print(f"\n{'='*70}")
|
|
print("Extracting REAL video features...")
|
|
print(f"{'='*70}")
|
|
|
|
real_features = load_or_compute_features(videos=real_videos,
|
|
extractor=extractor,
|
|
config=config,
|
|
cache_path=config.cache_real_features,
|
|
cache_name="real_features")
|
|
|
|
if verbose:
|
|
print(f"\n{'='*70}")
|
|
print("Extracting GENERATED video features...")
|
|
print(f"{'='*70}")
|
|
|
|
gen_features = load_or_compute_features(videos=gen_videos,
|
|
extractor=extractor,
|
|
config=config,
|
|
cache_path=None,
|
|
cache_name="gen_features")
|
|
|
|
if verbose:
|
|
print(f"\nReal videos/clips: {len(real_features)}")
|
|
print(f"Generated videos/clips: {len(gen_features)}")
|
|
print(f"\n{'='*70}")
|
|
print("Computing statistics...")
|
|
print(f"{'='*70}")
|
|
|
|
mu_real, sigma_real = compute_statistics(real_features)
|
|
mu_gen, sigma_gen = compute_statistics(gen_features)
|
|
|
|
if verbose:
|
|
print(f"\n{'='*70}")
|
|
print("Computing Fréchet distance...")
|
|
print(f"{'='*70}")
|
|
|
|
fvd = compute_frechet_distance(mu_real, sigma_real, mu_gen, sigma_gen)
|
|
|
|
if verbose:
|
|
print(f"\n{'='*70}")
|
|
print(f"FVD Score ({config.extractor_model.upper()}): {fvd:.4f}")
|
|
print(f"Protocol: {config}")
|
|
print(f"{'='*70}\n")
|
|
|
|
results = {
|
|
'fvd': fvd,
|
|
'protocol': str(config),
|
|
'model': config.extractor_model,
|
|
'config': config.to_dict(),
|
|
}
|
|
|
|
return results
|
|
|
|
|
|
def compute_fvd(real_videos: str | Path | torch.Tensor,
|
|
gen_videos: str | Path | torch.Tensor,
|
|
num_frames: int = 16,
|
|
batch_size: int = 32,
|
|
device: str = 'cuda',
|
|
num_videos: int | None = 2048,
|
|
cache_real_features: str | None = None,
|
|
i3d_model_path: str | None = None,
|
|
seed: int | None = None,
|
|
verbose: bool = True) -> float:
|
|
"""
|
|
Backward compatibility wrapper for computing FVD (defaults to I3D).
|
|
"""
|
|
num_videos = num_videos if num_videos is not None else 2048
|
|
|
|
config = FVDConfig(
|
|
num_videos=num_videos,
|
|
num_frames_per_clip=num_frames,
|
|
extractor_model='i3d', # Default to I3D
|
|
batch_size=batch_size,
|
|
device=device,
|
|
cache_real_features=cache_real_features,
|
|
i3d_model_path=i3d_model_path,
|
|
seed=seed,
|
|
)
|
|
|
|
result = compute_fvd_with_config(real_videos, gen_videos, config, verbose)
|
|
return result['fvd']
|