| 113 | artanh = lambda x: 0.5 * np.log((1 + x) / (1 - x)) |
| 114 | |
| 115 | class GSPMDecoder(nn.Module): |
| 116 | def __init__(self, |
| 117 | opt: Options, |
| 118 | transformer_dim: int, |
| 119 | mlp_dim=None, |
| 120 | init_density=0.2, |
| 121 | clip_scaling=0.1, |
| 122 | bias=True, |
| 123 | ): |
| 124 | super(GSPMDecoder, self).__init__() |
| 125 | |
| 126 | self.embed_dim = transformer_dim |
| 127 | if mlp_dim is not None: |
| 128 | self.mlp_dim = mlp_dim |
| 129 | else: |
| 130 | self.mlp_dim = transformer_dim |
| 131 | self.clip_scaling = clip_scaling |
| 132 | self.opt = opt |
| 133 | |
| 134 | token_w = self.opt.down_resolution[1] * (2 ** self.opt.decoder_ratio / self.opt.patch_size) # (512 / 2) |
| 135 | token_h = self.opt.down_resolution[0] * (2 ** self.opt.decoder_ratio / self.opt.patch_size) # (288 / 2) when decoder_ratio = 3 |
| 136 | w_coords = torch.linspace(-1, 1, int(token_w) + 1, dtype=torch.float32) |
| 137 | w_coords = (w_coords[1:] + w_coords[:-1]) / 2 |
| 138 | h_coords = torch.linspace(1, -1, int(token_h) + 1, dtype=torch.float32) |
| 139 | h_coords = (h_coords[1:] + h_coords[:-1]) / 2 |
| 140 | z_map, x_map = torch.meshgrid(h_coords, w_coords, indexing='ij') |
| 141 | self.register_buffer("z_map", z_map.flatten()) |
| 142 | self.register_buffer("x_map", x_map.flatten()) |
| 143 | |
| 144 | self.z_max, self.x_max = 1. / (token_h), 1. / (token_w) |
| 145 | |
| 146 | self.ratio = 1 |
| 147 | self.actual_h = int(self.opt.down_resolution[0] * (2 ** self.opt.decoder_ratio / self.opt.patch_size)) |
| 148 | self.actual_w = int(self.opt.down_resolution[1] * (2 ** self.opt.decoder_ratio / self.opt.patch_size)) |
| 149 | self.pixelshuffle = nn.PixelShuffle(upscale_factor=int(self.ratio**(0.5))) |
| 150 | |
| 151 | self.all_keys = ["xyz_static", "xyz_dynamic", |
| 152 | "rot_static", "rot_dynamic", |
| 153 | "opacity", "opacity_dynamic", |
| 154 | "scale", "rgb"] |
| 155 | self.key_dims = {"xyz_static": 3, "xyz_dynamic": 3 * (opt.forder), |
| 156 | "rot_static": 4, "rot_dynamic": 4, |
| 157 | "opacity": 1, "opacity_dynamic": 1, |
| 158 | "scale": 2, "rgb": 3, |
| 159 | "xz_scale": 2, "y_scale": 1} |
| 160 | |
| 161 | self.pred_keys = opt.pred_keys |
| 162 | self.has_pred = len(self.pred_keys) > 0 |
| 163 | |
| 164 | self.sample_keys = opt.sample_keys |
| 165 | |
| 166 | self.opacity_activation = opt.opacity_activation |
| 167 | |
| 168 | self.fix_keys = opt.fix_keys |
| 169 | |
| 170 | self.keep_dynamic = False # default to False, change in splatpredictor |
| 171 | |
| 172 | self.scale_min = 0.001 |