Spaces:
Sleeping
Sleeping
| import numpy as np | |
| from pathlib import Path | |
| from PIL import Image | |
| import json | |
| import torch | |
| from torch.utils.data import Dataset, DataLoader, default_collate | |
| from torchvision.transforms import ToTensor, Normalize, Compose, Resize | |
| from pytorch_lightning import LightningDataModule | |
| from einops import rearrange | |
| class LatentObjaverseSpiral(Dataset): | |
| def __init__( | |
| self, | |
| root_dir, | |
| split="train", | |
| transform=None, | |
| random_front=False, | |
| max_item=None, | |
| cond_aug_mean=-3.0, | |
| cond_aug_std=0.5, | |
| condition_on_elevation=False, | |
| **unused_kwargs, | |
| ): | |
| print("Using LVIS subset with precomputed Latents") | |
| self.root_dir = Path(root_dir) | |
| self.split = split | |
| self.random_front = random_front | |
| self.transform = transform | |
| self.latent_dir = Path("/mnt/vepfs/3Ddataset/render_results/latents512") | |
| self.ids = json.load(open("./assets/lvis_uids.json", "r")) | |
| self.n_views = 18 | |
| valid_ids = [] | |
| for idx in self.ids: | |
| if (self.root_dir / idx).exists(): | |
| valid_ids.append(idx) | |
| self.ids = valid_ids | |
| print("=" * 30) | |
| print("Number of valid ids: ", len(self.ids)) | |
| print("=" * 30) | |
| self.cond_aug_mean = cond_aug_mean | |
| self.cond_aug_std = cond_aug_std | |
| self.condition_on_elevation = condition_on_elevation | |
| if max_item is not None: | |
| self.ids = self.ids[:max_item] | |
| ## debug | |
| self.ids = self.ids * 10000 | |