| """AIGVDBench Open-Source split dataset (test only, frame-folder layout). |
| |
| Expected directory structure under ``data_root``:: |
| |
| Real/test/Real/videos/<vid>/frame_NNN.{jpg,png} (real) |
| OpenSource/<modality>/<gen>/test/videos/<vid>/frame_NNN.{jpg,png} (fake) |
| OpenSource/<modality>/<gen>/videos/<vid>/... (fallback for some gens) |
| |
| where ``<modality>`` is one of ``T2V``, ``I2V``, ``V2V`` and ``<gen>`` is the |
| generator name (e.g. ``Wan2.1``, ``CogVideoX``, ...). Each video is a folder |
| of frame images (one image per frame). |
| |
| Frames are uniformly sub-sampled to ``num_frames`` indices, center-cropped to |
| the shorter side, resized to ``crop_size``, and normalized with the dataset |
| mean/std (0.45 / 0.225; the eval loop later renormalizes to the backbone's |
| expected statistics). |
| """ |
| import glob |
| import os |
| from dataclasses import dataclass |
| from typing import List |
|
|
| import numpy as np |
| import torch |
| import torchvision.transforms.functional as TF |
| from PIL import Image |
|
|
|
|
| MEAN = (0.45, 0.45, 0.45) |
| STD = (0.225, 0.225, 0.225) |
| IMG_EXTS = ("*.png", "*.jpg", "*.jpeg") |
| TEST_MODALITIES = ("T2V", "I2V", "V2V") |
|
|
|
|
| @dataclass |
| class OpenSourceConfig: |
| data_root: str |
| num_frames: int = 8 |
| crop_size: int = 224 |
| generator_filter: List[str] = None |
| max_per_category: int = 0 |
|
|
|
|
| def _scan(data_root: str, generator_filter, max_per_cat: int): |
| """Return list of (video_dir, label, category_name) tuples for the test |
| set. Categories are namespaced by modality (e.g. ``T2V/Wan2.1``).""" |
| if not os.path.isdir(data_root): |
| raise RuntimeError(f"data_root not found: {data_root}") |
|
|
| out = [] |
|
|
| real_dir = os.path.join(data_root, "Real", "test", "Real", "videos") |
| if os.path.isdir(real_dir): |
| dirs = sorted(os.listdir(real_dir)) |
| if max_per_cat > 0: |
| dirs = dirs[:max_per_cat] |
| for v in dirs: |
| p = os.path.join(real_dir, v) |
| if os.path.isdir(p): |
| out.append((p, 0, "real")) |
|
|
| for modality in TEST_MODALITIES: |
| mod_dir = os.path.join(data_root, "OpenSource", modality) |
| if not os.path.isdir(mod_dir): |
| continue |
| for gen in sorted(os.listdir(mod_dir)): |
| cat_name = f"{modality}/{gen}" |
| if generator_filter and (gen not in generator_filter |
| and cat_name not in generator_filter): |
| continue |
| gen_dir = None |
| for candidate in ( |
| os.path.join(mod_dir, gen, "test", "videos"), |
| os.path.join(mod_dir, gen, "videos"), |
| ): |
| if os.path.isdir(candidate): |
| gen_dir = candidate |
| break |
| if gen_dir is None: |
| continue |
| dirs = sorted(os.listdir(gen_dir)) |
| if max_per_cat > 0: |
| dirs = dirs[:max_per_cat] |
| for v in dirs: |
| p = os.path.join(gen_dir, v) |
| if os.path.isdir(p): |
| out.append((p, 1, cat_name)) |
| return out |
|
|
|
|
| class OpenSourceTest(torch.utils.data.Dataset): |
| """Test-split frame loader. Returns ``(clip, label, idx, category)``.""" |
| def __init__(self, cfg: OpenSourceConfig): |
| self.cfg = cfg |
| self.samples = _scan(cfg.data_root, |
| cfg.generator_filter or [], |
| cfg.max_per_category) |
| if not self.samples: |
| raise RuntimeError(f"no test samples found under {cfg.data_root}") |
|
|
| def __len__(self): |
| return len(self.samples) |
|
|
| def _frames(self, video_dir): |
| paths = [] |
| for ext in IMG_EXTS: |
| paths.extend(glob.glob(os.path.join(video_dir, ext))) |
| return sorted(paths) |
|
|
| def _indices(self, total): |
| T = self.cfg.num_frames |
| if total < T: |
| return list(range(total)) + [total - 1] * (T - total) |
| return np.linspace(0, total - 1, T, dtype=int).tolist() |
|
|
| def __getitem__(self, idx): |
| video_dir, label, category = self.samples[idx] |
| paths = self._frames(video_dir) |
| if not paths: |
| raise RuntimeError(f"No frames in {video_dir}") |
| indices = self._indices(len(paths)) |
| imgs = [Image.open(paths[i]).convert("RGB") for i in indices] |
|
|
| W, H = imgs[0].size |
| s = min(W, H) |
| cs = self.cfg.crop_size |
| cx = (W - s) // 2 |
| cy = (H - s) // 2 |
|
|
| tensors = [] |
| for im in imgs: |
| im = TF.crop(im, cy, cx, s, s) |
| im = TF.resize(im, (cs, cs), interpolation=TF.InterpolationMode.BICUBIC) |
| t = TF.to_tensor(im) |
| t = TF.normalize(t, MEAN, STD) |
| tensors.append(t) |
| clip = torch.stack(tensors, dim=1) |
| return clip, label, idx, category |
|
|
|
|
| def collate(batch): |
| clips = torch.stack([b[0] for b in batch], dim=0) |
| labels = torch.tensor([b[1] for b in batch], dtype=torch.long) |
| cats = [b[3] for b in batch] |
| return clips, labels, cats |
|
|
