mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-13 03:14:00 +08:00
add name changes
This commit is contained in:
Justin committed
1 parent
d66a67fd8e
commit
29b7cb2b26
2 files changed
+265
No files matched your search
@@ -0,0 +1,124 @@
|
||||
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 PatchEmbeddingDataloader(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)
|
||||
@@ -0,0 +1,141 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from lerf.data.utils.feature_dataloader import FeatureDataloader
|
||||
from lerf.data.utils.patch_embedding_dataloader import PatchEmbeddingDataloader
|
||||
from lerf.encoders.image_encoder import BaseImageEncoder
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
class PyramidEmbeddingDataloader(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, PatchEmbeddingDataloader 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] = PatchEmbeddingDataloader(
|
||||
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, PatchEmbeddingDataloader 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
|
||||
Reference in new issue
Block a user