mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-13 03:14:00 +08:00
rename interpolator and patch dataloader
This commit is contained in:
Justin committed
1 parent
664e6e2874
commit
d66a67fd8e
4 files changed
+19
-289
No files matched your search
@@ -18,30 +18,31 @@ Datamanager.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os.path as osp
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple, Type, Union, Literal
|
||||
import os.path as osp
|
||||
import yaml
|
||||
from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union
|
||||
|
||||
import torch
|
||||
import yaml
|
||||
from rich.progress import Console
|
||||
|
||||
from nerfstudio.cameras.rays import RayBundle
|
||||
from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate
|
||||
from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes
|
||||
from nerfstudio.model_components.ray_generators import RayGenerator
|
||||
from nerfstudio.utils.misc import IterableWrapper
|
||||
from nerfstudio.cameras.rays import RayBundle
|
||||
|
||||
CONSOLE = Console(width=120)
|
||||
|
||||
from lerf.data.utils.dino_dataloader import DinoDataloader
|
||||
from lerf.data.utils.pyramid_embedding_dataloader import PyramidEmbeddingDataloader
|
||||
from lerf.encoders.image_encoder import BaseImageEncoder
|
||||
|
||||
from nerfstudio.data.datamanagers.base_datamanager import (
|
||||
VanillaDataManager,
|
||||
VanillaDataManagerConfig,
|
||||
)
|
||||
from lerf.data.utils.dino_dataloader import DinoDataloader
|
||||
from lerf.data.utils.patched_dataloader import PatchedDataloader
|
||||
from lerf.encoders.image_encoder import BaseImageEncoder
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -51,6 +52,7 @@ class LERFDataManagerConfig(VanillaDataManagerConfig):
|
||||
patch_tile_size_res: int = 7
|
||||
patch_stride_scaler: float = 0.5
|
||||
|
||||
|
||||
class LERFDataManager(VanillaDataManager): # pylint: disable=abstract-method
|
||||
"""Basic stored data manager implementation.
|
||||
|
||||
@@ -81,28 +83,28 @@ class LERFDataManager(VanillaDataManager): # pylint: disable=abstract-method
|
||||
self.image_encoder: BaseImageEncoder = kwargs["image_encoder"]
|
||||
images = [self.train_dataset[i]["image"].permute(2, 0, 1)[None, ...] for i in range(len(self.train_dataset))]
|
||||
images = torch.cat(images)
|
||||
|
||||
|
||||
cache_dir = f"outputs/{self.config.dataparser.data.name}"
|
||||
# clip_cache_path, dino_cache_path = self._search_cache(cache_dir)
|
||||
clip_cache_path = osp.join(cache_dir, "clip")
|
||||
dino_cache_path = osp.join(cache_dir, "dino.npy")
|
||||
self.clip_interpolator = PatchedDataloader(
|
||||
self.clip_interpolator = PyramidEmbeddingDataloader(
|
||||
image_list=images,
|
||||
device=self.device,
|
||||
cfg=
|
||||
{
|
||||
cfg={
|
||||
"tile_size_range": (0.05, 0.5),
|
||||
"tile_size_res": 7,
|
||||
"stride_scaler": 0.5,
|
||||
"image_shape": tuple(images.shape[2:4])
|
||||
"image_shape": tuple(images.shape[2:4]),
|
||||
},
|
||||
cache_path=clip_cache_path,
|
||||
model=self.image_encoder,
|
||||
)
|
||||
self.dino_dataloader = DinoDataloader(
|
||||
image_list=images, device=self.device,
|
||||
image_list=images,
|
||||
device=self.device,
|
||||
cfg={"image_shape": tuple(images.shape[2:4])},
|
||||
cache_path=dino_cache_path
|
||||
cache_path=dino_cache_path,
|
||||
)
|
||||
|
||||
# def _search_cache(self, cache_path, prop_dict):
|
||||
|
||||
@@ -1,148 +0,0 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
from lerf.data.utils.feature_dataloader import FeatureDataloader
|
||||
from lerf.data.utils.pyramid_interpolator import PyramidInterpolator
|
||||
from lerf.encoders.image_encoder import BaseImageEncoder
|
||||
|
||||
class PatchedDataloader(FeatureDataloader):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cfg: dict,
|
||||
device: torch.device,
|
||||
model: BaseImageEncoder,
|
||||
image_list: torch.Tensor = None,
|
||||
cache_path: str = None,
|
||||
):
|
||||
assert "tile_size_range" in cfg
|
||||
assert "tile_size_res" in cfg
|
||||
assert "stride_scaler" in cfg
|
||||
assert "image_shape" in cfg
|
||||
|
||||
cfg["tile_sizes"] = torch.linspace(*cfg["tile_size_range"], cfg["tile_size_res"]).to(device)
|
||||
cfg["strider_scaler_list"] = [
|
||||
self._stride_scaler(tr.item(), cfg["stride_scaler"]) for tr in cfg["tile_sizes"]
|
||||
]
|
||||
|
||||
self.model = model
|
||||
self.embed_size = self.model.embedding_dim
|
||||
self.data_dict = {}
|
||||
super().__init__(cfg, device, image_list, cache_path)
|
||||
|
||||
def __call__(self, img_points, scale=None):
|
||||
if scale is None:
|
||||
return self._random_scales(img_points)
|
||||
else:
|
||||
return self._uniform_scales(img_points, scale)
|
||||
|
||||
def _stride_scaler(self, tile_ratio, stride_scaler):
|
||||
return np.interp(tile_ratio, [0.05, 0.15], [1.0, stride_scaler])
|
||||
|
||||
def load(self):
|
||||
# don't create anything, PyramidInterpolator will create itself
|
||||
raise FileNotFoundError # trigger create
|
||||
|
||||
def create(self, image_list):
|
||||
os.makedirs(self.cache_path, exist_ok=True)
|
||||
for i, tr in enumerate(tqdm(self.cfg["tile_sizes"], desc="Scales")):
|
||||
stride_scaler = self.cfg["strider_scaler_list"][i]
|
||||
self.data_dict[i] = PyramidInterpolator(
|
||||
cfg={
|
||||
"tile_ratio": tr.item(),
|
||||
"stride_ratio": stride_scaler,
|
||||
"image_shape": self.cfg['image_shape']
|
||||
},
|
||||
device=self.device,
|
||||
model=self.model,
|
||||
image_list=image_list,
|
||||
cache_path=f'{self.cache_path}/level_{i}.npy',
|
||||
)
|
||||
|
||||
def save(self):
|
||||
# don't save anything, PyramidInterpolator will save itself
|
||||
pass
|
||||
|
||||
def _random_scales(self, img_points):
|
||||
# img_points: (B, 3) # (img_ind, x, y)
|
||||
# return: (B, 512), some random scale (between 0, 1)
|
||||
tile_sizes = self.cfg["tile_sizes"]
|
||||
|
||||
random_scale_bin = torch.randint(tile_sizes.shape[0] - 1, size=(img_points.shape[0],), device=self.device)
|
||||
random_scale_weight = torch.rand(img_points.shape[0], dtype=torch.float16, device=self.device)
|
||||
|
||||
stepsize = (tile_sizes[1] - tile_sizes[0]) / (tile_sizes[-1] - tile_sizes[0])
|
||||
|
||||
bottom_interp = torch.zeros((img_points.shape[0], self.embed_size), dtype=torch.float16, device=self.device)
|
||||
top_interp = torch.zeros((img_points.shape[0], self.embed_size), dtype=torch.float16, device=self.device)
|
||||
|
||||
for i in range(len(tile_sizes) - 1):
|
||||
bottom_interp[random_scale_bin == i] = self.data_dict[i](img_points[random_scale_bin == i])
|
||||
top_interp[random_scale_bin == i] = self.data_dict[i + 1](img_points[random_scale_bin == i])
|
||||
|
||||
return (
|
||||
torch.lerp(bottom_interp, top_interp, random_scale_weight[..., None]),
|
||||
(random_scale_bin * stepsize + random_scale_weight * stepsize)[..., None],
|
||||
)
|
||||
|
||||
def _uniform_scales(self, img_points, scale):
|
||||
tile_sizes = self.cfg["tile_sizes"]
|
||||
|
||||
scale_bin = torch.floor(
|
||||
(scale - tile_sizes[0]) / (tile_sizes[-1] - tile_sizes[0]) * (tile_sizes.shape[0] - 1)
|
||||
).to(torch.int64)
|
||||
scale_weight = (scale - tile_sizes[scale_bin]) / (tile_sizes[scale_bin + 1] - tile_sizes[scale_bin])
|
||||
interp_lst = torch.stack([interp(img_points) for interp in self.data_dict.values()])
|
||||
point_inds = torch.arange(img_points.shape[0])
|
||||
interp = torch.lerp(
|
||||
interp_lst[scale_bin, point_inds],
|
||||
interp_lst[scale_bin + 1, point_inds],
|
||||
torch.Tensor([scale_weight]).half().to(self.device)[..., None],
|
||||
)
|
||||
return interp / interp.norm(dim=-1, keepdim=True), scale
|
||||
|
||||
# def sample_pixels(self, N_patch: int, N_samp: int, interp=True):
|
||||
# remnant code for deblur LERF (could be interesting still)
|
||||
# """
|
||||
# returns {'indices':(N_patch*N_samp, 3), 'image',(N_patch*N_samp, 3),
|
||||
# 'clip': (N_patch, 512)}, clip_scale:((N_patch*N_samp, 1))
|
||||
# """
|
||||
# # 1. sample patch_ids: N_patch samples into the whole dataset. (N_patch,4) 4 == (im_id,x_ind,y_ind,scale)
|
||||
# im_ids = torch.randint(self.image_list.shape[0], (N_patch, 1))
|
||||
# col_ind = torch.randint(self.image_list.shape[3], (N_patch, 1))
|
||||
# row_ind = torch.randint(self.image_list.shape[2], (N_patch, 1))
|
||||
# scale_ind = torch.randint(len(self.data_dict), (N_patch, 1))
|
||||
# patch_ids = torch.concat([im_ids, row_ind, col_ind, scale_ind], dim=-1).to(self.device)
|
||||
# # 3. use the interpolators to upsample each patch sample into (N_patch*N_samp,4)
|
||||
# # 4 == (im_id,x_ind,y_ind,scale), and (N_patch,512)
|
||||
# # this is batch['indices'] (after removing scale) and batch['clip']
|
||||
# upsampled_ids, clips, scales = [], [], []
|
||||
# for i in self.data_dict:
|
||||
# relevant_ids = patch_ids[patch_ids[:, 3] == i, :3]
|
||||
# if interp:
|
||||
# res = self.data_dict[i].upsample_interp(relevant_ids, N_samp)
|
||||
# else:
|
||||
# res = self.data_dict[i].upsample(relevant_ids, N_samp)
|
||||
# if res is None:
|
||||
# continue
|
||||
# u, c = res
|
||||
# s = (
|
||||
# torch.rand((u.shape[0], 1), device=u.device, dtype=torch.float32) * 0.1
|
||||
# + self.cfg["tile_sizes"][i] / 2
|
||||
# - 0.05
|
||||
# )
|
||||
# upsampled_ids.append(u)
|
||||
# clips.append(c)
|
||||
# scales.append(s)
|
||||
# upsampled_ids, clips, scales = (
|
||||
# torch.concat(upsampled_ids, dim=0),
|
||||
# torch.concat(clips, dim=0),
|
||||
# torch.concat(scales, dim=0),
|
||||
# )
|
||||
# # 4. index into the images to get the color values (this is batch['image'])
|
||||
# image = self.image_list[upsampled_ids[:, 0], :, upsampled_ids[:, 1], upsampled_ids[:, 2]]
|
||||
# # 5. clip_scale is the 4th dim of step 3
|
||||
# return {"image": image, "clip": clips, "indices": upsampled_ids.cpu()}, scales
|
||||
@@ -1,124 +0,0 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
from lerf.data.utils.feature_dataloader import FeatureDataloader
|
||||
from lerf.encoders.image_encoder import BaseImageEncoder
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
class PyramidInterpolator(FeatureDataloader):
|
||||
def __init__(
|
||||
self,
|
||||
cfg: dict,
|
||||
device: torch.device,
|
||||
model: BaseImageEncoder,
|
||||
image_list: torch.Tensor = None,
|
||||
cache_path: str = None,
|
||||
):
|
||||
assert "tile_ratio" in cfg
|
||||
assert "stride_ratio" in cfg
|
||||
assert "image_shape" in cfg
|
||||
|
||||
self.kernel_size = int(cfg["image_shape"][0] * cfg["tile_ratio"])
|
||||
self.stride = int(self.kernel_size * cfg["stride_ratio"])
|
||||
self.padding = self.kernel_size // 2
|
||||
self.center_x = (
|
||||
(self.kernel_size - 1) / 2
|
||||
- self.padding
|
||||
+ self.stride
|
||||
* np.arange(
|
||||
np.floor((cfg["image_shape"][0] + 2 * self.padding - (self.kernel_size - 1) - 1) / self.stride + 1)
|
||||
)
|
||||
)
|
||||
self.center_y = (
|
||||
(self.kernel_size - 1) / 2
|
||||
- self.padding
|
||||
+ self.stride
|
||||
* np.arange(
|
||||
np.floor((cfg["image_shape"][1] + 2 * self.padding - (self.kernel_size - 1) - 1) / self.stride + 1)
|
||||
)
|
||||
)
|
||||
self.center_x = torch.from_numpy(self.center_x).to(device)
|
||||
self.center_y = torch.from_numpy(self.center_y).to(device)
|
||||
|
||||
self.model = model
|
||||
self.embed_size = self.model.embedding_dim
|
||||
super().__init__(cfg, device, image_list, cache_path)
|
||||
|
||||
def create(self, image_list):
|
||||
assert self.model is not None, "model must be provided to generate features"
|
||||
assert image_list is not None, "image_list must be provided to generate features"
|
||||
|
||||
unfold_func = torch.nn.Unfold(
|
||||
kernel_size=self.kernel_size,
|
||||
stride=self.stride,
|
||||
padding=self.padding,
|
||||
).to(self.device)
|
||||
|
||||
img_embeds = []
|
||||
for img in tqdm(image_list, desc="Embedding images", leave=False):
|
||||
img_embeds.append(self._embed_clip_tiles(img.unsqueeze(0), unfold_func))
|
||||
self.data = torch.from_numpy(np.stack(img_embeds)).half()
|
||||
|
||||
def __call__(self, img_points, mode="interp"):
|
||||
# img_points: (B, 3) # (img_ind, x, y) (img_ind, row, col)
|
||||
# return: (B, 512)
|
||||
img_points = img_points.to(self.device)
|
||||
img_ind, img_points_x, img_points_y = img_points[:, 0], img_points[:, 1], img_points[:, 2]
|
||||
|
||||
x_ind = torch.searchsorted(self.center_x, img_points_x, side="left") - 1
|
||||
y_ind = torch.searchsorted(self.center_y, img_points_y, side="left") - 1
|
||||
return self._interp_inds(img_ind, x_ind, y_ind, img_points_x, img_points_y)
|
||||
|
||||
def _interp_inds(self, img_ind, x_ind, y_ind, img_points_x, img_points_y):
|
||||
topleft = self.data[img_ind, x_ind, y_ind].to(self.device)
|
||||
topright = self.data[img_ind, x_ind + 1, y_ind].to(self.device)
|
||||
botleft = self.data[img_ind, x_ind, y_ind + 1].to(self.device)
|
||||
botright = self.data[img_ind, x_ind + 1, y_ind + 1].to(self.device)
|
||||
|
||||
x_stride = self.stride
|
||||
y_stride = self.stride
|
||||
right_w = ((img_points_x - (self.center_x[x_ind])) / x_stride).half()
|
||||
top = torch.lerp(topleft, topright, right_w[:, None])
|
||||
bot = torch.lerp(botleft, botright, right_w[:, None])
|
||||
|
||||
bot_w = ((img_points_y - (self.center_y[y_ind])) / y_stride).half()
|
||||
return torch.lerp(top, bot, bot_w[:, None])
|
||||
|
||||
def _embed_clip_tiles(self, image, unfold_func):
|
||||
# image augmentation: slow-ish (0.02s for 600x800 image per augmentation)
|
||||
aug_imgs = torch.cat([image])
|
||||
|
||||
tiles = unfold_func(aug_imgs).permute(2, 0, 1).reshape(-1, 3, self.kernel_size, self.kernel_size).to("cuda")
|
||||
|
||||
with torch.no_grad():
|
||||
clip_embeds = self.model.encode_image(tiles)
|
||||
clip_embeds /= clip_embeds.norm(dim=-1, keepdim=True)
|
||||
|
||||
clip_embeds = clip_embeds.reshape((self.center_x.shape[0], self.center_y.shape[0], -1))
|
||||
clip_embeds = torch.concat((clip_embeds, clip_embeds[:, [-1], :]), dim=1)
|
||||
clip_embeds = torch.concat((clip_embeds, clip_embeds[[-1], :, :]), dim=0)
|
||||
return clip_embeds.detach().cpu().numpy()
|
||||
|
||||
# def upsample_interp(self, input_ids, N_samp):
|
||||
# B = input_ids.shape[0]
|
||||
# if B == 0:
|
||||
# return None
|
||||
# img_ind, img_points_x, img_points_y = input_ids[:, 0], input_ids[:, 1], input_ids[:, 2]
|
||||
|
||||
# x_ind = torch.searchsorted(self.center_x, img_points_x, side="left") - 1
|
||||
# y_ind = torch.searchsorted(self.center_y, img_points_y, side="left") - 1
|
||||
# gt_embeds = self.interp_inds(img_ind, x_ind, y_ind, img_points_x, img_points_y)
|
||||
|
||||
# blur_size = 0.5 * self.kernel_size
|
||||
|
||||
# x_start = img_points_x - blur_size / 2 # (N_patch, 1)
|
||||
# x_start = x_start[:, None].repeat_interleave(N_samp, 0) # (N_patcH*N_samp,1)
|
||||
# y_start = img_points_y - blur_size / 2
|
||||
# y_start = y_start[:, None].repeat_interleave(N_samp, 0)
|
||||
|
||||
# new_xs = torch.rand((N_samp * B, 1), device=self.device) * blur_size + x_start
|
||||
# new_ys = torch.rand((N_samp * B, 1), device=self.device) * blur_size + y_start
|
||||
# new_xs = torch.clip(new_xs, 0, self.imlist_shape[2] - 1)
|
||||
# new_ys = torch.clip(new_ys, 0, self.imlist_shape[3] - 1)
|
||||
# new_imids = img_ind[:, None].repeat_interleave(N_samp, 0)
|
||||
# return torch.concat((new_imids, new_xs.long(), new_ys.long()), dim=1), gt_embeds.to(self.device)
|
||||
@@ -37,7 +37,7 @@ lerf_method = MethodSpecification(
|
||||
),
|
||||
),
|
||||
model=LERFModelConfig(
|
||||
eval_num_rays_per_chunk=1 << 16,
|
||||
eval_num_rays_per_chunk=1 << 15,
|
||||
hashgrid_sizes=(19, 19),
|
||||
hashgrid_layers=(12, 12),
|
||||
hashgrid_resolutions=((16, 128), (128, 512)),
|
||||
@@ -83,7 +83,7 @@ lerf_method_big = MethodSpecification(
|
||||
),
|
||||
),
|
||||
model=LERFModelConfig(
|
||||
eval_num_rays_per_chunk=1 << 16,
|
||||
eval_num_rays_per_chunk=1 << 15,
|
||||
hashgrid_sizes=(19, 19),
|
||||
hashgrid_layers=(14, 14),
|
||||
hashgrid_resolutions=((16, 128), (128, 512)),
|
||||
@@ -130,7 +130,7 @@ lerf_method_lite = MethodSpecification(
|
||||
),
|
||||
),
|
||||
model=LERFModelConfig(
|
||||
eval_num_rays_per_chunk=1 << 16,
|
||||
eval_num_rays_per_chunk=1 << 15,
|
||||
hashgrid_sizes=(19,),
|
||||
hashgrid_layers=(16,),
|
||||
hashgrid_resolutions=((16, 512),),
|
||||
|
||||
Reference in new issue
Block a user