| 165 | |
| 166 | |
| 167 | class PointNet_SA_Module(nn.Module): |
| 168 | def __init__(self, npoint, nsample, radius, in_channel, mlp, if_bn=True, group_all=False, use_xyz=True): |
| 169 | """ |
| 170 | Args: |
| 171 | npoint: int, number of points to sample |
| 172 | nsample: int, number of points in each local region |
| 173 | radius: float |
| 174 | in_channel: int, input channel of features(points) |
| 175 | mlp: list of int, |
| 176 | """ |
| 177 | super(PointNet_SA_Module, self).__init__() |
| 178 | self.npoint = npoint |
| 179 | self.nsample = nsample |
| 180 | self.radius = radius |
| 181 | self.mlp = mlp |
| 182 | self.group_all = group_all |
| 183 | self.use_xyz = use_xyz |
| 184 | if use_xyz: |
| 185 | in_channel += 3 |
| 186 | |
| 187 | last_channel = in_channel |
| 188 | self.mlp_conv = [] |
| 189 | for out_channel in mlp: |
| 190 | self.mlp_conv.append(Conv2d(last_channel, out_channel, if_bn=if_bn)) |
| 191 | last_channel = out_channel |
| 192 | |
| 193 | self.mlp_conv = nn.Sequential(*self.mlp_conv) |
| 194 | |
| 195 | def forward(self, xyz, points): |
| 196 | """ |
| 197 | Args: |
| 198 | xyz: Tensor, (B, 3, N) |
| 199 | points: Tensor, (B, f, N) |
| 200 | |
| 201 | Returns: |
| 202 | new_xyz: Tensor, (B, 3, npoint) |
| 203 | new_points: Tensor, (B, mlp[-1], npoint) |
| 204 | """ |
| 205 | if self.group_all: |
| 206 | new_xyz, new_points, idx, grouped_xyz = sample_and_group_all(xyz, points, self.use_xyz) |
| 207 | else: |
| 208 | new_xyz, new_points, idx, grouped_xyz = sample_and_group(xyz, points, self.npoint, self.nsample, self.radius, self.use_xyz) |
| 209 | |
| 210 | new_points = self.mlp_conv(new_points) |
| 211 | new_points = torch.max(new_points, 3)[0] |
| 212 | |
| 213 | return new_xyz, new_points |
| 214 | |
| 215 | |
| 216 | class PointNet_FP_Module(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected