diff --git a/projects-lerf/lerf/data/lerf_datamanager.py b/projects-lerf/lerf/data/lerf_datamanager.py index 43609e754..da627eb86 100644 --- a/projects-lerf/lerf/data/lerf_datamanager.py +++ b/projects-lerf/lerf/data/lerf_datamanager.py @@ -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): diff --git a/projects-lerf/lerf/data/utils/patched_dataloader.py b/projects-lerf/lerf/data/utils/patched_dataloader.py deleted file mode 100644 index 43998432e..000000000 --- a/projects-lerf/lerf/data/utils/patched_dataloader.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/projects-lerf/lerf/data/utils/pyramid_interpolator.py b/projects-lerf/lerf/data/utils/pyramid_interpolator.py deleted file mode 100644 index 7a95d9188..000000000 --- a/projects-lerf/lerf/data/utils/pyramid_interpolator.py +++ /dev/null @@ -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) diff --git a/projects-lerf/lerf/lerf_config.py b/projects-lerf/lerf/lerf_config.py index 3b4997b0e..ceaf6f4e5 100644 --- a/projects-lerf/lerf/lerf_config.py +++ b/projects-lerf/lerf/lerf_config.py @@ -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),),