293 lines
8.6 KiB
Python
293 lines
8.6 KiB
Python
import dataclasses
|
|
import importlib
|
|
import math
|
|
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 Bool, 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)
|
|
|
|
|
|
def tri_winding(tri: Float[Tensor, "*B 3 2"]) -> Float[Tensor, "*B 3 3"]:
|
|
# One pad for determinant
|
|
tri_sq = F.pad(tri, (0, 1), "constant", 1.0)
|
|
det_tri = torch.det(tri_sq)
|
|
tri_rev = torch.cat(
|
|
(tri_sq[..., 0:1, :], tri_sq[..., 2:3, :], tri_sq[..., 1:2, :]), -2
|
|
)
|
|
tri_sq[det_tri < 0] = tri_rev[det_tri < 0]
|
|
return tri_sq
|
|
|
|
|
|
def triangle_intersection_2d(
|
|
t1: Float[Tensor, "*B 3 2"],
|
|
t2: Float[Tensor, "*B 3 2"],
|
|
eps=1e-12,
|
|
) -> Float[Tensor, "*B"]: # noqa: F821
|
|
"""Returns True if triangles collide, False otherwise"""
|
|
|
|
def chk_edge(x: Float[Tensor, "*B 3 3"]) -> Bool[Tensor, "*B"]: # noqa: F821
|
|
logdetx = torch.logdet(x.double())
|
|
if eps is None:
|
|
return ~torch.isfinite(logdetx)
|
|
return ~(torch.isfinite(logdetx) & (logdetx > math.log(eps)))
|
|
|
|
t1s = tri_winding(t1)
|
|
t2s = tri_winding(t2)
|
|
|
|
# Assume the triangles do not collide in the begging
|
|
ret = torch.zeros(t1.shape[0], dtype=torch.bool, device=t1.device)
|
|
for i in range(3):
|
|
edge = torch.roll(t1s, i, dims=1)[:, :2, :]
|
|
# Check if all points of triangle 2 lay on the external side of edge E.
|
|
# If this is the case the triangle do not collide
|
|
upd = (
|
|
chk_edge(torch.cat((edge, t2s[:, 0:1]), 1))
|
|
& chk_edge(torch.cat((edge, t2s[:, 1:2]), 1))
|
|
& chk_edge(torch.cat((edge, t2s[:, 2:3]), 1))
|
|
)
|
|
# Here no collision is still True due to inversion
|
|
ret = ret | upd
|
|
|
|
for i in range(3):
|
|
edge = torch.roll(t2s, i, dims=1)[:, :2, :]
|
|
|
|
upd = (
|
|
chk_edge(torch.cat((edge, t1s[:, 0:1]), 1))
|
|
& chk_edge(torch.cat((edge, t1s[:, 1:2]), 1))
|
|
& chk_edge(torch.cat((edge, t1s[:, 2:3]), 1))
|
|
)
|
|
# Here no collision is still True due to inversion
|
|
ret = ret | upd
|
|
|
|
return ~ret # Do the inversion
|
|
|
|
|
|
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
|
|
|
|
if 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)
|