import dataclasses import importlib from dataclasses import dataclass from typing import Any, List, Optional, Tuple, Union import numpy as np import PIL import torch import torch.nn as nn import torch.nn.functional as F from jaxtyping import Float, Int, Num from omegaconf import DictConfig, OmegaConf from torch import Tensor class BaseModule(nn.Module): @dataclass class Config: pass cfg: Config # add this to every subclass of BaseModule to enable static type checking def __init__( self, cfg: Optional[Union[dict, DictConfig]] = None, *args, **kwargs ) -> None: super().__init__() self.cfg = parse_structured(self.Config, cfg) self.configure(*args, **kwargs) def configure(self, *args, **kwargs) -> None: raise NotImplementedError def find_class(cls_string): module_string = ".".join(cls_string.split(".")[:-1]) cls_name = cls_string.split(".")[-1] module = importlib.import_module(module_string, package=None) cls = getattr(module, cls_name) return cls def parse_structured(fields: Any, cfg: Optional[Union[dict, DictConfig]] = None) -> Any: # Check if cfg.keys are in fields cfg_ = cfg.copy() keys = list(cfg_.keys()) field_names = {f.name for f in dataclasses.fields(fields)} for key in keys: # This is helpful when swapping out modules from CLI if key not in field_names: print(f"Ignoring {key} as it's not supported by {fields}") cfg_.pop(key) scfg = OmegaConf.merge(OmegaConf.structured(fields), cfg_) return scfg EPS_DTYPE = { torch.float16: 1e-4, torch.bfloat16: 1e-4, torch.float32: 1e-7, torch.float64: 1e-8, } def dot(x, y, dim=-1): return torch.sum(x * y, dim, keepdim=True) def reflect(x, n): return x - 2 * dot(x, n) * n def normalize(x, dim=-1, eps=None): if eps is None: eps = EPS_DTYPE[x.dtype] return F.normalize(x, dim=dim, p=2, eps=eps) ValidScale = Union[Tuple[float, float], Num[Tensor, "2 D"]] def scale_tensor( dat: Num[Tensor, "... D"], inp_scale: ValidScale, tgt_scale: ValidScale ): if inp_scale is None: inp_scale = (0, 1) if tgt_scale is None: tgt_scale = (0, 1) if isinstance(tgt_scale, Tensor): assert dat.shape[-1] == tgt_scale.shape[-1] dat = (dat - inp_scale[0]) / (inp_scale[1] - inp_scale[0]) dat = dat * (tgt_scale[1] - tgt_scale[0]) + tgt_scale[0] return dat def dilate_fill(img, mask, iterations=10): oldMask = mask.float() oldImg = img mask_kernel = torch.ones( (1, 1, 3, 3), dtype=oldMask.dtype, device=oldMask.device, ) for i in range(iterations): newMask = torch.nn.functional.max_pool2d(oldMask, 3, 1, 1) # Fill the extension with mean color of old valid regions img_unfold = F.unfold(oldImg, (3, 3)).view(1, 3, 3 * 3, -1) mask_unfold = F.unfold(oldMask, (3, 3)).view(1, 1, 3 * 3, -1) new_mask_unfold = F.unfold(newMask, (3, 3)).view(1, 1, 3 * 3, -1) # Average color of the valid region mean_color = (img_unfold.sum(dim=2) / mask_unfold.sum(dim=2).clip(1)).unsqueeze( 2 ) # Extend it to the new region fill_color = (mean_color * new_mask_unfold).view(1, 3 * 3 * 3, -1) mask_conv = F.conv2d( newMask, mask_kernel, padding=1 ) # Get the sum for each kernel patch newImg = F.fold( fill_color, (img.shape[-2], img.shape[-1]), (3, 3) ) / mask_conv.clamp(1) diffMask = newMask - oldMask oldMask = newMask oldImg = torch.lerp(oldImg, newImg, diffMask) return oldImg def float32_to_uint8_np( x: Float[np.ndarray, "*B H W C"], dither: bool = True, dither_mask: Optional[Float[np.ndarray, "*B H W C"]] = None, dither_strength: float = 1.0, ) -> Int[np.ndarray, "*B H W C"]: if dither: dither = ( dither_strength * np.random.rand(*x[..., :1].shape).astype(np.float32) - 0.5 ) if dither_mask is not None: dither = dither * dither_mask return np.clip(np.floor((256.0 * x + dither)), 0, 255).astype(np.uint8) return np.clip(np.floor((256.0 * x)), 0, 255).astype(torch.uint8) def convert_data(data): if data is None: return None elif isinstance(data, np.ndarray): return data elif isinstance(data, torch.Tensor): if data.dtype in [torch.float16, torch.bfloat16]: data = data.float() return data.detach().cpu().numpy() elif isinstance(data, list): return [convert_data(d) for d in data] elif isinstance(data, dict): return {k: convert_data(v) for k, v in data.items()} else: raise TypeError( "Data must be in type numpy.ndarray, torch.Tensor, list or dict, getting", type(data), ) class ImageProcessor: def convert_and_resize( self, image: Union[PIL.Image.Image, np.ndarray, torch.Tensor], size: int, ): if isinstance(image, PIL.Image.Image): image = torch.from_numpy(np.array(image).astype(np.float32) / 255.0) elif isinstance(image, np.ndarray): if image.dtype == np.uint8: image = torch.from_numpy(image.astype(np.float32) / 255.0) else: image = torch.from_numpy(image) elif isinstance(image, torch.Tensor): pass batched = image.ndim >= 4 view_batch = image.ndim >= 5 if view_batch: image = image.view(-1, *image.shape[2:]) elif not batched: image = image[None, ...] image = F.interpolate( image.permute(0, 3, 1, 2), (size, size), mode="bilinear", align_corners=False, antialias=True, ).permute(0, 2, 3, 1) if not batched: image = image[0] return image def __call__( self, image: Union[ PIL.Image.Image, np.ndarray, torch.FloatTensor, List[PIL.Image.Image], List[np.ndarray], List[torch.FloatTensor], ], size: int, ) -> Any: if isinstance(image, (np.ndarray, torch.FloatTensor)) and image.ndim == 4: image = self.convert_and_resize(image, size) else: if not isinstance(image, list): image = [image] image = [self.convert_and_resize(im, size) for im in image] image = torch.stack(image, dim=0) return image def get_intrinsic_from_fov(fov, H, W, bs=-1): focal_length = 0.5 * H / np.tan(0.5 * fov) intrinsic = np.identity(3, dtype=np.float32) intrinsic[0, 0] = focal_length intrinsic[1, 1] = focal_length intrinsic[0, 2] = W / 2.0 intrinsic[1, 2] = H / 2.0 if bs > 0: intrinsic = intrinsic[None].repeat(bs, axis=0) return torch.from_numpy(intrinsic)