def __init__(
self,
img: DistanceImage,
segments: int,
spline: int = 5,
seg_emb_size: int = 128,
anchor_interim: int = 128,
):
super().__init__()
self.img = img
self.spline_n = spline
self.ses = ses = seg_emb_size
segments = int(segments)
self.segments = segments
self.midsize = anchor_interim
self.seg_emb = nn.Embedding(segments + 2, ses)
self.anchors = SkipConnMLP(
in_size=2,
out=2 + anchor_interim,
latent_size=ses,
num_layers=5,
hidden_size=512,
init="xavier",
)
self.point_estim = SkipConnMLP(
in_size=2,
out=(spline - 2) * 2,
num_layers=5,
hidden_size=512,
latent_size=ses + 2 * anchor_interim,
init="xavier",
)
def __init__(
self,
latent_size=32,
):
super().__init__()
self.ls = latent_size
self.approx = SkipConnMLP(
in_size=2, out=3, latent_size=latent_size,
enc=FourierEncoder(input_dims=2, freqs=32),
num_layers=5, hidden_size=256,
init="siren", activation=torch.sin,
)
self.displacement = SkipConnMLP(
in_size=3, out=1, latent_size=latent_size,
hidden_size=256, activation=torch.sin,
)
def __init__(
self,
output_latent_size: int = 0,
bounds: float = 1.5,
# this latent size is per resolution
latent_size: int = 32,
# number of different resolutions to try the SDF at.
resolutions: int = 3,
):
assert (resolutions > 0), "Must have at least 1 SDF resolution"
super().__init__()
self.latent_size = resolutions * latent_size
self.res = resolutions
self.tiers = nn.ModuleList([
SkipConnMLP(
in_size=3,
out=1 + output_latent_size,
latent_size=latent_size,
enc=FourierEncoder(input_dims=3),
activation=torch.sin,
num_layers=4,
hidden_size=256,
skip=3,
init="xavier",
) for _ in range(resolutions)
])
self.assigned_latent = None
self.bounds = bounds
def __init__(self, img, *args, **kwargs):
super().__init__()
self.img = img
self.pred = SkipConnMLP(
in_size=1,
out=2,
num_layers=7,
hidden_size=512,
init="xavier",
)
def __init__(self):
super().__init__()
self.query = SkipConnMLP(
in_size=2,
out=1,
latent_size=0,
activation=torch.sin,
num_layers=5,
hidden_size=256,
init="siren",
)
def __init__(self, n=32):
super().__init__()
self.n = n
self.points = nn.Parameter(torch.randn(n, 2, requires_grad=True),
requires_grad=True)
self.query = SkipConnMLP(
in_size=n,
out=1,
num_layers=5,
hidden_size=512,
init="xavier",
)
def __init__(self, reso: int = 16, emb_size: int = 128):
super().__init__()
self.grid = nn.Parameter(torch.randn(1, emb_size, reso, reso))
self.query = SkipConnMLP(
in_size=emb_size,
out=1,
latent_size=2,
activation=torch.sin,
num_layers=5,
hidden_size=256,
init="siren",
)
class SmoothImageApprox(nn.Module):
def __init__(
self,
latent_size=32,
):
super().__init__()
self.ls = latent_size
self.approx = SkipConnMLP(
in_size=2, out=3, latent_size=latent_size,
enc=FourierEncoder(input_dims=2, freqs=32),
num_layers=5, hidden_size=256,
init="siren", activation=torch.sin,
)
self.displacement = SkipConnMLP(
in_size=3, out=1, latent_size=latent_size,
hidden_size=256, activation=torch.sin,
)
# initializes this approximation at 0 to be the correct image
def init_zero(self, image, epochs=750):
sz = image.shape[0]
device = image.device
points = torch.stack(torch.meshgrid(
torch.linspace(-1, 1, steps=sz, device=device),
torch.linspace(-1, 1, steps=sz, device=device),
), dim=-1)
opt = optim.Adam(self.approx.parameters(), lr=1e-3)
sched = optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs, eta_min=5e-5)
t = trange(epochs)
for i in t:
opt.zero_grad()
Z = torch.randn(sz,sz,self.ls,device=device)
got = wide_sigmoid(self.approx(points, Z))
loss = F.mse_loss(got, image).sqrt()
loss.backward()
t.set_postfix(l2=f"{loss.item():.03f}")
opt.step()
sched.step()
return got.detach()
def forward(self, x, t, latent):
angle = self.displacement(torch.cat([x, t], dim=-1), latent)
angle = torch.remainder(angle, 2*math.pi)
dx = torch.cat([angle.cos(), angle.sin()], dim=-1)
# Force there to be some movement
dx = F.normalize(dx, dim=-1, eps=1e-6) * (t * 2 * math.pi).sin()
return wide_sigmoid(self.approx(torch.fmod(x + dx, 1), latent))
def __init__(
self,
latent_size: int = 32,
output_latent_size: int = 0,
bounds: float = 1.5,
):
super().__init__()
self.latent_size = latent_size
self.assigned_latent = None
self.mlp = SkipConnMLP(
in_size=3,
out=1 + output_latent_size,
latent_size=latent_size,
enc=FourierEncoder(input_dims=3),
activation=torch.sin,
num_layers=7,
hidden_size=512,
skip=3,
init="siren",
)
self.bounds = bounds