在点云上进行深度学习：在Google Colab中实现PointNet

import numpy as npimport randomimport math!pip install path.py;from path import Path

!wget http://3dvision.princeton.edu/projects/2014/3DShapeNets/ModelNet10.zip!unzip -q ModelNet10.zip path = Path("ModelNet10")

def read_off(file):    if 'OFF' != file.readline().strip():        raise('Not a valid OFF header')    n_verts, n_faces, __ = tuple([int(s) for s in file.readline().strip().split(' ')])    verts = [[float(s) for s in file.readline().strip().split(' ')] for i_vert in range(n_verts)]    faces = [[int(s) for s in file.readline().strip().split(' ')][1:] for i_face in range(n_faces)]    return verts, faces    with open(path/"bed/train/bed_0001.off", 'r') as f:    mesh = read_off(f)

verts, faces = meshareas = np.zeros((len(faces)))verts = np.array(verts) # function to calculate triangle area by its vertices# https://en.wikipedia.org/wiki/Heron%27s_formuladef triangle_area(pt1, pt2, pt3):    side_a = np.linalg.norm(pt1 - pt2)    side_b = np.linalg.norm(pt2 - pt3)    side_c = np.linalg.norm(pt3 - pt1)    s = 0.5 * ( side_a + side_b + side_c)    return max(s * (s - side_a) * (s - side_b) * (s - side_c), 0)**0.5 # we calculate areas of all faces in our meshfor i in range(len(areas)):    areas[i] = (triangle_area(verts[faces[i][0]],                              verts[faces[i][1]],                              verts[faces[i][2]]))

k = 3000# we sample 'k' faces with probabilities proportional to their areas# weights are used to create a distribution.# they don't have to sum up to one.sampled_faces = (random.choices(faces,                                weights=areas,                                k=k)) # function to sample points on a triangle surfacedef sample_point(pt1, pt2, pt3):    # barycentric coordinates on a triangle    # https://mathworld.wolfram.com/BarycentricCoordinates.html    s, t = sorted([random.random(), random.random()])    f = lambda i: s * pt1[i] + (t-s) * pt2[i] + (1-t) * pt3[i]    return (f(0), f(1), f(2)) pointcloud = np.zeros((k, 3)) # sample points on chosen faces for the point cloud of size 'k'for i in range(len(sampled_faces)):    pointcloud[i] = (sample_point(verts[sampled_faces[i][0]],                                  verts[sampled_faces[i][1]],                                  verts[sampled_faces[i][2]]))

# normalizenorm_pointcloud = pointcloud - np.mean(pointcloud, axis=0)norm_pointcloud /= np.max(np.linalg.norm(norm_pointcloud, axis=1)) # rotation around z-axistheta = random.random() * 2. * math.pi # rotation anglerot_matrix = np.array([[ math.cos(theta), -math.sin(theta),    0],                       [ math.sin(theta),  math.cos(theta),    0],                       [0,                             0,      1]]) rot_pointcloud = rot_matrix.dot(pointcloud.T).T # add some noisenoise = np.random.normal(0, 0.02, (pointcloud.shape))noisy_pointcloud = rot_pointcloud + noise

import torchimport torch.nn as nnimport torch.nn.functional as F class Tnet(nn.Module):   def __init__(self, k=3):      super().__init__()      self.k=k      self.conv1 = nn.Conv1d(k,64,1)      self.conv2 = nn.Conv1d(64,128,1)      self.conv3 = 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.bn1 = nn.BatchNorm1d(64)      self.bn2 = nn.BatchNorm1d(128)      self.bn3 = nn.BatchNorm1d(1024)      self.bn4 = nn.BatchNorm1d(512)      self.bn5 = nn.BatchNorm1d(256)           def forward(self, input):      # input.shape == (bs,n,3)      bs = input.size(0)      xb = F.relu(self.bn1(self.conv1(input)))      xb = F.relu(self.bn2(self.conv2(xb)))      xb = F.relu(self.bn3(self.conv3(xb)))      pool = nn.MaxPool1d(xb.size(-1))(xb)      flat = nn.Flatten(1)(pool)      xb = F.relu(self.bn4(self.fc1(flat)))      xb = F.relu(self.bn5(self.fc2(xb)))            # initialize as identity      init = torch.eye(self.k, requires_grad=True).repeat(bs,1,1)      if xb.is_cuda:        init=init.cuda()      # add identity to the output      matrix = self.fc3(xb).view(-1,self.k,self.k) + init      return matrix

class Transform(nn.Module):   def __init__(self):        super().__init__()        self.input_transform = Tnet(k=3)        self.feature_transform = Tnet(k=64)        self.conv1 = nn.Conv1d(3,64,1)         self.conv2 = nn.Conv1d(64,128,1)        self.conv3 = nn.Conv1d(128,1024,1)         self.bn1 = nn.BatchNorm1d(64)        self.bn2 = nn.BatchNorm1d(128)        self.bn3 = nn.BatchNorm1d(1024)          def forward(self, input):        matrix3x3 = self.input_transform(input)        # batch matrix multiplication        xb = torch.bmm(torch.transpose(input,1,2), matrix3x3).transpose(1,2)        xb = F.relu(self.bn1(self.conv1(xb)))         matrix64x64 = self.feature_transform(xb)        xb = torch.bmm(torch.transpose(xb,1,2), matrix64x64).transpose(1,2)         xb = F.relu(self.bn2(self.conv2(xb)))        xb = self.bn3(self.conv3(xb))        xb = nn.MaxPool1d(xb.size(-1))(xb)        output = nn.Flatten(1)(xb)        return output, matrix3x3, matrix64x64

class PointNet(nn.Module):    def __init__(self, classes=10):        super().__init__()        self.transform = Transform()        self.fc1 = nn.Linear(1024, 512)        self.fc2 = nn.Linear(512, 256)        self.fc3 = nn.Linear(256, classes)         self.bn1 = nn.BatchNorm1d(512)        self.bn2 = nn.BatchNorm1d(256)        self.dropout = nn.Dropout(p=0.3)        self.logsoftmax = nn.LogSoftmax(dim=1)     def forward(self, input):        xb, matrix3x3, matrix64x64 = self.transform(input)        xb = F.relu(self.bn1(self.fc1(xb)))        xb = F.relu(self.bn2(self.dropout(self.fc2(xb))))        output = self.fc3(xb)        return self.logsoftmax(output), matrix3x3, matrix64x64

def pointnetloss(outputs, labels, m3x3, m64x64, alpha = 0.0001):    criterion = torch.nn.NLLLoss()    bs = outputs.size(0)    id3x3 = torch.eye(3, requires_grad=True).repeat(bs, 1, 1)    id64x64 = torch.eye(64, requires_grad=True).repeat(bs, 1, 1)    if outputs.is_cuda:        id3x3 = id3x3.cuda()        id64x64 = id64x64.cuda()    diff3x3 = id3x3 - torch.bmm(m3x3, m3x3.transpose(1, 2))    diff64x64 = id64x64 - torch.bmm(m64x64, m64x64.transpose(1, 2))    return criterion(outputs, labels) + alpha * (torch.norm(diff3x3) + torch.norm(diff64x64)) / float(bs)