Spaces:
Running on Zero
Running on Zero
File size: 4,439 Bytes
e793773 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """DL3DV loader.
On-disk layout::
<data_root>/<resolution>/<scene>/
transforms.json frames[] (OpenGL c2w) + fl_x/w/h
wan2_caption.json {"caption": ...}
images_4/<name>.png downsampled RGB frames
DL3DV's ``transforms.json`` uses the Nerfstudio/OpenGL convention (camera looks
along -Z). We convert each c2w to OpenCV (``c2w @ diag(1, -1, -1, 1)``) at the
loader boundary so downstream Plucker rays point the correct way. DL3DV clips
are pre-cut, so a fixed stride of 1 is used.
"""
from __future__ import annotations
import json
import random
from pathlib import Path
from typing import Any
import numpy as np
from PIL import Image
from scope.data.common import BaseCameraVideoDataset, linspace_indices
_GL2CV = np.diag([1.0, -1.0, -1.0, 1.0]).astype(np.float32)
class DL3DVDataset(BaseCameraVideoDataset):
def __init__(
self,
data_root: str,
num_frames: int = 81,
height: int = 480,
width: int = 832,
sample_stride: int = 1,
random_sample: bool = True,
max_videos: int | None = None,
near_depth_map: dict[str, float] | None = None,
trajectory_scale: float = 1.0,
return_first_frame: bool = True,
) -> None:
super().__init__(
num_frames=num_frames,
height=height,
width=width,
near_depth_map=near_depth_map,
trajectory_scale=trajectory_scale,
return_first_frame=return_first_frame,
)
self.data_root = Path(data_root)
self.sample_stride = sample_stride
self.random_sample = random_sample
scenes = self._scan_scenes()
kept = self.filter_ids_by_near_depth([s.name for s in scenes])
keep = set(kept)
self.scenes = [s for s in scenes if s.name in keep]
if max_videos is not None:
self.scenes = self.scenes[:max_videos]
if not self.scenes:
raise ValueError(f"No DL3DV scenes found under {self.data_root}")
print(f"[DL3DV] {len(self.scenes)} scenes")
def _scan_scenes(self) -> list[Path]:
index_path = self.data_root / "valid_video_dirs.json"
if index_path.exists():
return [Path(p) for p in json.loads(index_path.read_text(encoding="utf-8"))]
scenes: list[Path] = []
for res_dir in sorted(self.data_root.iterdir()):
if not res_dir.is_dir():
continue
for scene in sorted(res_dir.iterdir()):
if (
(scene / "transforms.json").exists()
and (scene / "wan2_caption.json").exists()
and (scene / "images_4").exists()
):
scenes.append(scene)
try:
index_path.write_text(json.dumps([str(p) for p in scenes]), encoding="utf-8")
except OSError:
pass
return scenes
def __len__(self) -> int:
return len(self.scenes)
def _sample_indices(self, total: int) -> list[int]:
stride = self.sample_stride
if not self.random_sample or total < self.num_frames * stride:
return linspace_indices(total, self.num_frames)
max_start = max(0, total - (self.num_frames - 1) * stride - 1)
start = random.randint(0, max_start) if max_start > 0 else 0
return [start + i * stride for i in range(self.num_frames)]
def _load_raw(self, index: int) -> dict[str, Any]:
scene = self.scenes[index]
meta = json.loads((scene / "transforms.json").read_text(encoding="utf-8"))
frames = meta["frames"]
indices = self._sample_indices(len(frames))
pil_frames = [
Image.open(scene / "images_4" / Path(frames[i]["file_path"]).name).convert("RGB")
for i in indices
]
poses = np.stack(
[
(np.asarray(frames[i]["transform_matrix"], dtype=np.float32) @ _GL2CV)[:3]
for i in indices
]
)
caption = json.loads((scene / "wan2_caption.json").read_text(encoding="utf-8")).get(
"caption", ""
)
x_fov = 2.0 * np.arctan(meta["w"] / (2.0 * meta["fl_x"]))
return {
"video_id": scene.name,
"frames": pil_frames,
"poses": poses,
"x_fov": float(x_fov),
"xi": 0.0,
"caption": caption,
}
|