hofee/nbv_reconstruction
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update modules and pipeline

Browse Source
This commit is contained in:
hofee
2024-08-21 17:57:52 +08:00
parent 913d4e521d
commit 837e1c870a
8 changed files with 129 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

89
modules/pts_encoder/pointnet_encoder.py
View File

@@ -9,57 +9,21 @@ import torch.nn.functional as F
from modules.pts_encoder.abstract_pts_encoder import PointsEncoder
import PytorchBoot.stereotype as stereotype
class STNkd(nn.Module):
def __init__(self, k=64):
super(STNkd, self).__init__()
self.conv1 = torch.nn.Conv1d(k, 64, 1)
self.conv2 = torch.nn.Conv1d(64, 128, 1)
self.conv3 = torch.nn.Conv1d(128, 1024, 1)
self.fc1 = nn.Linear(1024, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, k * k)
self.relu = nn.ReLU()
self.k = k
def forward(self, x):
batchsize = x.size()[0]
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = F.relu(self.conv3(x))
x = torch.max(x, 2, keepdim=True)[0]
x = x.view(-1, 1024)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
iden = (
Variable(torch.from_numpy(np.eye(self.k).flatten().astype(np.float32)))
.view(1, self.k * self.k)
.repeat(batchsize, 1)
)
if x.is_cuda:
iden = iden.to(x.get_device())
x = x + iden
x = x.view(-1, self.k, self.k)
return x
@stereotype.module("pointnet_encoder")
class PointNetEncoder(PointsEncoder):
def __init__(self, global_feat=True, in_dim=3, out_dim=1024, feature_transform=False):
def __init__(self, config:dict):
super(PointNetEncoder, self).__init__()
self.out_dim = out_dim
self.feature_transform = feature_transform
self.stn = STNkd(k=in_dim)
self.conv1 = torch.nn.Conv1d(in_dim, 64, 1)
self.out_dim = config["out_dim"]
self.in_dim = config["in_dim"]
self.feature_transform = config.get("feature_transform", False)
self.stn = STNkd(k=self.in_dim)
self.conv1 = torch.nn.Conv1d(self.in_dim , 64, 1)
self.conv2 = torch.nn.Conv1d(64, 128, 1)
self.conv3 = torch.nn.Conv1d(128, 512, 1)
self.conv4 = torch.nn.Conv1d(512, out_dim, 1)
self.global_feat = global_feat
self.conv4 = torch.nn.Conv1d(512, self.out_dim , 1)
self.global_feat = config["global_feat"]
if self.feature_transform:
self.f_stn = STNkd(k=64)
@@ -97,6 +61,41 @@ class PointNetEncoder(PointsEncoder):
pts_feature = self(pts)
return pts_feature
class STNkd(nn.Module):
def __init__(self, k=64):
super(STNkd, self).__init__()
self.conv1 = torch.nn.Conv1d(k, 64, 1)
self.conv2 = torch.nn.Conv1d(64, 128, 1)
self.conv3 = torch.nn.Conv1d(128, 1024, 1)
self.fc1 = nn.Linear(1024, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, k * k)
self.relu = nn.ReLU()
self.k = k
def forward(self, x):
batchsize = x.size()[0]
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = F.relu(self.conv3(x))
x = torch.max(x, 2, keepdim=True)[0]
x = x.view(-1, 1024)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
iden = (
Variable(torch.from_numpy(np.eye(self.k).flatten().astype(np.float32)))
.view(1, self.k * self.k)
.repeat(batchsize, 1)
)
if x.is_cuda:
iden = iden.to(x.get_device())
x = x + iden
x = x.view(-1, self.k, self.k)
return x
if __name__ == "__main__":
sim_data = Variable(torch.rand(32, 2500, 3))