File size: 4,699 Bytes
4fe19fe | 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 | import os
import torch
import numpy as np
from torch.utils.data import Dataset
from typing import Optional
from data.smal_quadrupeds import SMALQuadrupedGenerator, FAMILIES, PARENT_INDICES, GENERATOR_VERSION
from data.ot_coupling import precompute_ot_noise
class CreatureDataset(Dataset):
"""PyTorch Dataset for 3D creatures conditioned on kinematic skeletons."""
def __init__(
self,
num_samples_per_family: int = 50,
num_points: int = 2048,
augment: bool = True,
cache_dir: Optional[str] = "data/cache",
seed: int = 42,
ot_noise_per_sample: int = 0,
independent_noise_prob: float = 0.0,
ot_workers: Optional[int] = None
):
super().__init__()
self.num_points = num_points
self.augment = augment
self.cache_dir = cache_dir
self.independent_noise_prob = independent_noise_prob
self.samples = []
self.ot_noise = None
generator = SMALQuadrupedGenerator(seed=seed)
# Cache file is keyed on generation parameters so a config change never reuses stale samples
cache_file = os.path.join(cache_dir, f"dataset_v{GENERATOR_VERSION}_n{num_samples_per_family}_p{num_points}_s{seed}.pt") if cache_dir else None
if cache_file and os.path.exists(cache_file):
cached = torch.load(cache_file, weights_only=False)
self.samples = cached["samples"]
else:
for family in FAMILIES:
for _ in range(num_samples_per_family):
instance = generator.generate_instance(family=family, num_points=num_points)
self.samples.append(instance)
if cache_file:
os.makedirs(cache_dir, exist_ok=True)
torch.save({"samples": self.samples}, cache_file)
if ot_noise_per_sample > 0:
self.ot_noise = self._load_or_compute_ot_noise(cache_file, ot_noise_per_sample, seed, ot_workers)
def _load_or_compute_ot_noise(self, cache_file: Optional[str], num_noise: int, seed: int, workers: Optional[int]):
noise_file = cache_file.replace(".pt", f"_ot{num_noise}.pt") if cache_file else None
if noise_file and os.path.exists(noise_file):
return torch.load(noise_file, weights_only=True)["noise"]
noise = precompute_ot_noise([s["points"] for s in self.samples], num_noise, seed=seed, num_workers=workers)
noise = [torch.from_numpy(n) for n in noise]
if noise_file:
torch.save({"noise": noise}, noise_file)
return noise
def __len__(self) -> int:
return len(self.samples)
def _apply_augmentation(self, points: np.ndarray, joints: np.ndarray, noise: Optional[np.ndarray] = None):
# Coherent rotation around vertical Y axis
theta = np.random.uniform(-np.pi, np.pi)
cos_t, sin_t = np.cos(theta), np.sin(theta)
R = np.array([
[cos_t, 0.0, sin_t],
[0.0, 1.0, 0.0 ],
[-sin_t, 0.0, cos_t]
], dtype=np.float32)
points = points @ R
joints = joints @ R
# N(0, I) is rotation invariant: rotating the matched noise keeps it Gaussian and keeps the OT pairing
if noise is not None:
noise = noise @ R
# Scaling jitter
scale = np.random.uniform(0.92, 1.08)
points *= scale
joints *= scale
return points, joints, noise
def __getitem__(self, idx: int):
raw = self.samples[idx]
pts = raw["points"].copy()
jts = raw["joints"].copy()
noise = None
if self.ot_noise is not None:
if np.random.rand() < self.independent_noise_prob:
# Hybrid coupling: some independent noise keeps the field easy to learn near t = 0
noise = np.random.randn(*pts.shape).astype(np.float32)
else:
stack = self.ot_noise[idx]
noise = stack[np.random.randint(len(stack))].numpy().copy()
if self.augment:
pts, jts, noise = self._apply_augmentation(pts, jts, noise)
item = {
"points": torch.from_numpy(pts).float(),
"joints": torch.from_numpy(jts).float(),
"joint_parents": torch.from_numpy(raw["joint_parents"]).long(),
"family_id": torch.tensor(raw["family_id"]).long()
}
if noise is not None:
item["noise"] = torch.from_numpy(noise).float()
return item
|