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