| 144 | |
| 145 | |
| 146 | class PointPatchEmbed(nn.Module): |
| 147 | |
| 148 | def __init__(self, |
| 149 | sample_ratio=0.0625, |
| 150 | sample_number=1024, |
| 151 | group_size=32, |
| 152 | in_channels=6, |
| 153 | channels=1024, |
| 154 | kernel_size=1, |
| 155 | stride=1, |
| 156 | normalize_dp=False, |
| 157 | relative_xyz=True, |
| 158 | ): |
| 159 | super().__init__() |
| 160 | self.sample_ratio = sample_ratio |
| 161 | self.sample_number = sample_number |
| 162 | self.group_size = group_size |
| 163 | |
| 164 | self.sample_fn = furthest_point_sample |
| 165 | self.grouper = KNNGroup(self.group_size, relative_xyz=relative_xyz, normalize_dp=normalize_dp) |
| 166 | |
| 167 | self.conv1 = nn.Conv2d(in_channels, channels, kernel_size=kernel_size, stride=stride) |
| 168 | |
| 169 | |
| 170 | def forward(self, x): |
| 171 | # coordinates |
| 172 | p = x[:, :, 3:].contiguous() |
| 173 | |
| 174 | B, N, _ = p.shape[:3] |
| 175 | # idx = self.sample_fn(p, int(N * self.sample_ratio)).long() |
| 176 | idx = self.sample_fn(p, self.sample_number).long() |
| 177 | center_p = torch.gather(p, 1, idx.unsqueeze(-1).expand(-1, -1, 3)) |
| 178 | # query neighbors. |
| 179 | _, fj = self.grouper(center_p, p, x.permute(0, 2, 1).contiguous()) # [B, N, 6] -> [B, 6, N] -> [B, 6, 1024, 32] |
| 180 | |
| 181 | # [B, 6, 1024] -> [B, channels, 1024, 1] |
| 182 | fj = self.conv1(fj).max(dim=-1, keepdim=True)[0] |
| 183 | |
| 184 | return fj |
| 185 | |
| 186 | |
| 187 | if __name__ == '__main__': |
no outgoing calls
no test coverage detected