PointNet代码学习(pytorch版本)

源码地址

pointnet.pytorch
感谢大神!

代码结构

(pytorch) s@s:~/pointnet.pytorch$ tree -d
.
├── misc
├── pointnet
│   └── __pycache__
├── scripts
├── shapenetcore_partanno_segmentation_benchmark_v0
│   ├── 02691156
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   ├── 02773838
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   ├── 02954340
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   ├── 02958343
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   ├── 03001627
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   ├── 03261776
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   ├── 03467517
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   ├── 03624134
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   ├── 03636649
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   ├── 03642806
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   ├── 03790512
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   ├── 03797390
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   ├── 03948459
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   ├── 04099429
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   ├── 04225987
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   ├── 04379243
│   │   ├── points
│   │   ├── points_label
│   │   └── seg_img
│   └── train_test_split
└── utils
    ├── cls
    ├── __pycache__
    └── seg

74 directories

utils

├── cls
│   ├── cls_model_0.pth
│   └── cls_model_1.pth
├── point_test.pts
├── __pycache__
│   ├── show3d_balls.cpython-36.pyc
│   ├── show3d_balls.cpython-37.pyc
│   └── show_seg.cpython-36.pyc
├── render_balls_so.cpp
├── render_balls_so.so
├── seg
│   ├── seg_model_Chair_0.pth
│   ├── seg_model_Chair_1.pth
│   ├── seg_model_Chair_2.pth
│   ├── seg_model_Chair_3.pth
│   └── seg_model_Chair_4.pth
├── show3d_balls.py
├── show_cls.py
├── show_points.py
├── show_seg.py
├── train_classification.py
└── train_segmentation.py

3 directories, 19 files

简介

• cls和seg文件夹下的是模型;
• train_classification.py和train_segmentation.py是训练脚本
• show_seg.py和 show_cls.py是导入模型进行测试并可视化的脚本
• show3d_balls.py是可视化脚本,含有可视化相关的函数
• show_points.py是自己写的测试的脚本(可以无视)
• render_balls_so.cpp和render_balls_so.so是可视化渲染相关的库
• point_test.pts是点云文件,自己测试用(不要理会)

代码详细注释

• show_seg.py

'''
对原始点云进行分割,并可视化
例:python show_seg.py --model seg/seg_model_Chair_1.pth 
						--class_choice Airplane --idx 2
'''

from __future__ import print_function
from show3d_balls import showpoints
import argparse
import numpy as np
import torch
import torch.nn.parallel
import torch.utils.data
from torch.autograd import Variable

#先把系统文件夹调到pointnet.pytorch下,防止找不到pointnet这个文件夹,
import sys
sys.path.append('/home/s/pointnet.pytorch')

from pointnet.dataset import ShapeNetDataset
from pointnet.model import PointNetDenseCls
import matplotlib.pyplot as plt




# 命令行解析
parser = argparse.ArgumentParser()

parser.add_argument('--model', type=str, default='', help='model path')
parser.add_argument('--idx', type=int, default=0, help='model index')
parser.add_argument('--dataset', type=str, default='', help='dataset path')
parser.add_argument('--class_choice', type=str, default='', help='class choice')


# 输出一行状态栏参数如下:
# 				Namespace(class_choice='Airplane', dataset='', 
# 							idx=2, model='seg/seg_model_Chair_1.pth')
opt = parser.parse_args()		
# print("opt:{}".format(opt))

# 数据预处理,得到某一类模型的集合
d = ShapeNetDataset(
#    root=opt.dataset,
	root='/home/s/pointnet.pytorch/shapenetcore_partanno_segmentation_benchmark_v0',
    class_choice=[opt.class_choice],  #选择哪一类模型
    split='test',
    data_augmentation=0)

idx = opt.idx
# print("d:{}".format(d))
print("model %d/%d" % (idx, len(d)))  #d代表全部的飞机数量
# model 2/341

# print('dir(d):{}'.format(dir(d)))
# print('d[idx]:{}'.format(d[idx]))


point, seg = d[idx]  #模型里的第idx的点云
print(point.size(), seg.size())     #seg代表每一个点的标签
# torch.Size([2500, 3]) torch.Size([2500])

point_np = point.numpy() #将torch转为numpy

# 可视化
cmap = plt.cm.get_cmap("hsv", 10)
cmap = np.array([cmap(i) for i in range(10)])[:, :3]
gt = cmap[seg.numpy() - 1, :]

# 载入模型
state_dict = torch.load(opt.model)  
classifier = PointNetDenseCls(k= state_dict['conv4.weight'].size()[0])
classifier.load_state_dict(state_dict)
classifier.eval()  #评估

# 点云转置
point = point.transpose(1, 0).contiguous()
print('point.transpose(1, 0).shape: ',point.shape)

point = Variable(point.view(1, point.size()[0], point.size()[1]))
print('--------------------')
print(point.dtype)
pred, _, _ = classifier(point)   #分割
# print('\npred.shape:',pred[0],'\n')
pred_choice = pred.data.max(2)[1]
print(pred_choice.numpy())   #输出每一个点的预测类别

# print(pred_choice.size())
print(pred_choice.numpy()[0])  #[1 1 1 ... 1 1 1]
pred_color = cmap[pred_choice.numpy()[0], :]   #根据分类结果显示颜色
print('\npred_color: ',pred_color.shape,'\n')


showpoints(point_np, gt, pred_color)  #pred_colord的为(2500, 3)的矩阵
print(point_np.shape)

• show_cls.py

from __future__ import print_function  #使用python3的print函数
import argparse
import torch
import torch.nn.parallel
import torch.utils.data
from torch.autograd import Variable

import sys
sys.path.append('/home/s/pointnet.pytorch')

from pointnet.dataset import ShapeNetDataset  #ShapeNetDataset是一个类 下面会实例化这个类
from pointnet.model import PointNetCls   #读入模型
import torch.nn.functional as F


#showpoints(np.random.randn(2500,3), c1 = np.random.uniform(0,1,size = (2500)))

parser = argparse.ArgumentParser()

parser.add_argument('--model', type=str, default = '',  help='model path')
parser.add_argument('--num_points', type=int, default=2500, help='input batch size')


opt = parser.parse_args()
print(opt)


# 实例化ShapeNetDataset类
test_dataset = ShapeNetDataset(
#    root='shapenetcore_partanno_segmentation_benchmark_v0',
	root='/home/s/pointnet.pytorch/shapenetcore_partanno_segmentation_benchmark_v0',
    split='test',
    classification=True,
    npoints=opt.num_points,
    data_augmentation=False)

# 读入测试数据
testdataloader = torch.utils.data.DataLoader(
    test_dataset, batch_size=32, shuffle=True)


# 导入模型
classifier = PointNetCls(k=len(test_dataset.classes))
classifier.cuda()
classifier.load_state_dict(torch.load(opt.model))
classifier.eval()


for i, data in enumerate(testdataloader, 0):
    points, target = data
    points, target = Variable(points), Variable(target[:, 0])
    points = points.transpose(2, 1)
    points, target = points.cuda(), target.cuda()
    pred, _, _ = classifier(points)  #进行分类
    loss = F.nll_loss(pred, target)  #计算损失函数

# 计算准确率
    pred_choice = pred.data.max(1)[1]
    correct = pred_choice.eq(target.data).cpu().sum()
    print('i:%d  loss: %f accuracy: %f' % (i, loss.data.item(), correct / float(32))) 

• show_points.py

'''
自己写的,用来测试
可视化文件夹下的点云数据
输入:n*3的矩阵
'''

from __future__ import print_function
from show3d_balls import showpoints
import argparse
import numpy as np
import torch
import torch.nn.parallel
import torch.utils.data
from torch.autograd import Variable
from pointnet.model import PointNetDenseCls

# with open('./point_test.pts') as file:
#     for line in file:
#         print(len(line))

import matplotlib.pylab as plt
import sys

# from utils.show_seg import seg

sys.path.append('/home/s/pointnet.pytorch')

# points=np.loadtxt('./point_test.pts')
points=np.loadtxt('./point_test.pts',dtype=np.float32)  #预测只能输入float32的格式的数据

print(points.shape)

# 可视化
cmap = plt.cm.get_cmap("hsv", 10)
cmap = np.array([cmap(i) for i in range(10)])[:, :3]
# gt = cmap[seg.numpy() - 1, :]

# 可视化点云
showpoints(points)

#采样到2500个点
choice = np.random.choice(len(points), 2500, replace=True)
# print('choice:{}'.format(choice))
points = points[choice, :]
print('points[choice, :]:{}'.format(points))
point_np=points



# 载入模型
state_dict = torch.load('./seg/seg_model_Chair_1.pth')
classifier = PointNetDenseCls(k= state_dict['conv4.weight'].size()[0])
classifier.load_state_dict(state_dict)
classifier.eval()  #设置为评估状态

# 点云转置

points=torch.from_numpy(points)
print(points.shape)
point = points.transpose(1, 0).contiguous()
print('point.transpose(1, 0).shape: ',point.shape)

point = Variable(point.view(1, point.size()[0], point.size()[1]))  #转为torch变量1,3,2500
print('--------------------')

print(point.dtype)
# point=torch.tensor(point,dtype=torch.float32)
pred, _, _ = classifier(point)   #分割
print(pred)

pred_choice = pred.data.max(2)[1]
print(pred_choice.numpy())   #输出每一个点的预测类别

# print(pred_choice.size())
print(pred_choice.numpy()[0])  #[1 1 1 ... 1 1 1]
pred_color = cmap[pred_choice.numpy()[0], :]   #根据分类结果显示颜色
print('\npred_color: ',pred_color.shape,'\n')
print(pred_color.dtype)

# point_np=point.numpy().reshape(2500,3)
print(point_np.shape)
print(point_np.dtype)
showpoints(point_np, pred_color,pred_color)  #pred_colord的为(2500, 3)的矩阵

• show3d_balls.py

'''
点云数据可视化
'''

import numpy as np
import ctypes as ct
import cv2
import sys
showsz = 800
mousex, mousey = 0.5, 0.5
zoom = 1.0
changed = True

def onmouse(*args):
    global mousex, mousey, changed
    y = args[1]
    x = args[2]
    mousex = x / float(showsz)
    mousey = y / float(showsz)
    changed = True

cv2.namedWindow('show3d')
cv2.moveWindow('show3d', 0, 0)
cv2.setMouseCallback('show3d', onmouse)

dll = np.ctypeslib.load_library('render_balls_so', '.')

# 该函数输入为n*3的矩阵
def showpoints(xyz,c_gt=None, c_pred = None, waittime=0, 
    showrot=False, magnifyBlue=0, freezerot=False, background=(0,0,0), 
    normalizecolor=True, ballradius=10):
    global showsz, mousex, mousey, zoom, changed
    xyz=xyz-xyz.mean(axis=0)
    radius=((xyz**2).sum(axis=-1)**0.5).max()
    xyz/=(radius*2.2)/showsz
    if c_gt is None:
        c0 = np.zeros((len(xyz), ), dtype='float32') + 255
        c1 = np.zeros((len(xyz), ), dtype='float32') + 255
        c2 = np.zeros((len(xyz), ), dtype='float32') + 255
    else:
        c0 = c_gt[:, 0]
        c1 = c_gt[:, 1]
        c2 = c_gt[:, 2]


    if normalizecolor:
        c0 /= (c0.max() + 1e-14) / 255.0
        c1 /= (c1.max() + 1e-14) / 255.0
        c2 /= (c2.max() + 1e-14) / 255.0


    c0 = np.require(c0, 'float32', 'C')
    c1 = np.require(c1, 'float32', 'C')
    c2 = np.require(c2, 'float32', 'C')

    show = np.zeros((showsz, showsz, 3), dtype='uint8')
    def render():
        rotmat=np.eye(3)
        if not freezerot:
            xangle=(mousey-0.5)*np.pi*1.2
        else:
            xangle=0
        rotmat = rotmat.dot(
            np.array([
                [1.0, 0.0, 0.0],
                [0.0, np.cos(xangle), -np.sin(xangle)],
                [0.0, np.sin(xangle), np.cos(xangle)],
            ]))
        if not freezerot:
            yangle = (mousex - 0.5) * np.pi * 1.2
        else:
            yangle = 0
        rotmat = rotmat.dot(
            np.array([
                [np.cos(yangle), 0.0, -np.sin(yangle)],
                [0.0, 1.0, 0.0],
                [np.sin(yangle), 0.0, np.cos(yangle)],
            ]))
        rotmat *= zoom
        nxyz = xyz.dot(rotmat) + [showsz / 2, showsz / 2, 0]

        ixyz = nxyz.astype('int32')
        show[:] = background
        dll.render_ball(
            ct.c_int(show.shape[0]), ct.c_int(show.shape[1]),
            show.ctypes.data_as(ct.c_void_p), ct.c_int(ixyz.shape[0]),
            ixyz.ctypes.data_as(ct.c_void_p), c0.ctypes.data_as(ct.c_void_p),
            c1.ctypes.data_as(ct.c_void_p), c2.ctypes.data_as(ct.c_void_p),
            ct.c_int(ballradius))

        if magnifyBlue > 0:
            show[:, :, 0] = np.maximum(show[:, :, 0], np.roll(
                show[:, :, 0], 1, axis=0))
            if magnifyBlue >= 2:
                show[:, :, 0] = np.maximum(show[:, :, 0],
                                           np.roll(show[:, :, 0], -1, axis=0))
            show[:, :, 0] = np.maximum(show[:, :, 0], np.roll(
                show[:, :, 0], 1, axis=1))
            if magnifyBlue >= 2:
                show[:, :, 0] = np.maximum(show[:, :, 0],
                                           np.roll(show[:, :, 0], -1, axis=1))
        if showrot:
            cv2.putText(show, 'xangle %d' % (int(xangle / np.pi * 180)),
                        (30, showsz - 30), 0, 0.5, cv2.cv.CV_RGB(255, 0, 0))
            cv2.putText(show, 'yangle %d' % (int(yangle / np.pi * 180)),
                        (30, showsz - 50), 0, 0.5, cv2.cv.CV_RGB(255, 0, 0))
            cv2.putText(show, 'zoom %d%%' % (int(zoom * 100)), (30, showsz - 70), 0,
                        0.5, cv2.cv.CV_RGB(255, 0, 0))
    changed = True
    while True:
        if changed:
            render()
            changed = False
        cv2.imshow('show3d', show)
        if waittime == 0:
            cmd = cv2.waitKey(10) % 256
        else:
            cmd = cv2.waitKey(waittime) % 256
        if cmd == ord('q'):
            break
        elif cmd == ord('Q'):
            sys.exit(0)

        if cmd == ord('t') or cmd == ord('p'):
            if cmd == ord('t'):
                if c_gt is None:
                    c0 = np.zeros((len(xyz), ), dtype='float32') + 255
                    c1 = np.zeros((len(xyz), ), dtype='float32') + 255
                    c2 = np.zeros((len(xyz), ), dtype='float32') + 255
                else:
                    c0 = c_gt[:, 0]
                    c1 = c_gt[:, 1]
                    c2 = c_gt[:, 2]
            else:
                if c_pred is None:
                    c0 = np.zeros((len(xyz), ), dtype='float32') + 255
                    c1 = np.zeros((len(xyz), ), dtype='float32') + 255
                    c2 = np.zeros((len(xyz), ), dtype='float32') + 255
                else:
                    c0 = c_pred[:, 0]
                    c1 = c_pred[:, 1]
                    c2 = c_pred[:, 2]
            if normalizecolor:
                c0 /= (c0.max() + 1e-14) / 255.0
                c1 /= (c1.max() + 1e-14) / 255.0
                c2 /= (c2.max() + 1e-14) / 255.0
            c0 = np.require(c0, 'float32', 'C')
            c1 = np.require(c1, 'float32', 'C')
            c2 = np.require(c2, 'float32', 'C')
            changed = True

        if cmd==ord('n'):
            zoom*=1.1
            changed=True
        elif cmd==ord('m'):
            zoom/=1.1
            changed=True
        elif cmd==ord('r'):
            zoom=1.0
            changed=True
        elif cmd==ord('s'):
            cv2.imwrite('show3d.png',show)
        if waittime!=0:
            break
    return cmd

if __name__ == '__main__':
    np.random.seed(100)
    showpoints(np.random.randn(2500, 3))

• train_classification.py

• train_segmentation.py

'''
训练分割的网络模型
'''
# from __future__ import print_function
import argparse
import os
import random
import torch
import torch.nn.parallel
import torch.optim as optim
import torch.utils.data

import sys
sys.path.append('/home/s/pointnet.pytorch')

from pointnet.dataset import ShapeNetDataset
from pointnet.model import PointNetDenseCls, feature_transform_regularizer
import torch.nn.functional as F
from tqdm import tqdm
import numpy as np


parser = argparse.ArgumentParser()
parser.add_argument(
    '--batchSize', type=int, default=32, help='input batch size')#每一次输入的32个元素,训练一遍需要输入 Size/batchSize 次 
parser.add_argument(
    '--workers', type=int, help='number of data loading workers', default=4)
parser.add_argument(
    '--nepoch', type=int, default=25, help='number of epochs to train for')#对整个数据集训练25次
parser.add_argument('--outf', type=str, default='seg', help='output folder')
parser.add_argument('--model', type=str, default='', help='model path')
parser.add_argument('--dataset', type=str, required=True, help="dataset path")
parser.add_argument('--class_choice', type=str, default='Chair', help="class_choice")
parser.add_argument('--feature_transform', action='store_true', help="use feature transform")

opt = parser.parse_args()
# print(opt)

#使产生的随机数是确定的 保证结果是可以复现的
opt.manualSeed = random.randint(1, 10000)  # fix seed 产生一个随机点
# print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)   #本函数设置随机数的类型  没有返回值
torch.manual_seed(opt.manualSeed)

dataset = ShapeNetDataset(
    root=opt.dataset,
    classification=False,
    class_choice=[opt.class_choice])
dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=opt.batchSize,
    shuffle=True,
    num_workers=int(opt.workers))

test_dataset = ShapeNetDataset(
    root=opt.dataset,
    classification=False,
    class_choice=[opt.class_choice],
    split='test',
    data_augmentation=False)
testdataloader = torch.utils.data.DataLoader(
    test_dataset,
    batch_size=opt.batchSize,
    shuffle=True,
    num_workers=int(opt.workers))

# print(len(dataset), len(test_dataset))  #1958  341
num_classes = dataset.num_seg_classes   #num_seg_classes是一个数字,代表该类物体应该被分成几类,由class_choice在终端输入后根据字典找出来的
# print('-----classes', num_classes)  # 4
try:
    os.makedirs(opt.outf)
except OSError:
    pass

blue = lambda x: '\033[94m' + x + '\033[0m'  #设置显示的颜色

classifier = PointNetDenseCls(k=num_classes, feature_transform=opt.feature_transform)   #分割网络

if opt.model != '':   #如果模型存在就导入
    classifier.load_state_dict(torch.load(opt.model))

optimizer = optim.Adam(classifier.parameters(), lr=0.001, betas=(0.9, 0.999))
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)  # scheduler.step()调用step_size次,学习率才会调整一次
classifier.cuda()

num_batch = len(dataset) / opt.batchSize   #计算需要几个batch才能导入dataset

for epoch in range(opt.nepoch):
    # scheduler.step()   #更新一下学习率 每step_size=20调整一次
    # print('----------lr:{}'.format(classifier.optimizer.state_dict()['param_groups'][0]['lr'] ) )
    for i, data in enumerate(dataloader, 0):  #i 代表第几个batch(最大值:num_batch),每个batch有batcSize个点云,每个点云2500个点

        # print('data:{}'.format(data))

        points, target = data
        # print('points.shape:{}'.format(points.size))
        # print('target.shape:{}'.format(target.size))

        points = points.transpose(2, 1)
        points, target = points.cuda(), target.cuda()
        optimizer.zero_grad()   #梯度归零
        classifier = classifier.train()  #模型设置为训练模式
        pred, trans, trans_feat = classifier(points)
        pred = pred.view(-1, num_classes)   #num_classes代表应该被分割成多少类  torch.Size([20000, 4])
        target = target.view(-1, 1)[:, 0] - 1

        # print(pred.size(), target.size())  #输出pred和target的格式  torch.Size([20000, 4]) torch.Size([20000])


        loss = F.nll_loss(pred, target)  #计算损失
        if opt.feature_transform:
            loss += feature_transform_regularizer(trans_feat) * 0.001
        loss.backward()  #反向传播损失
        optimizer.step()    #梯度下降优化 以batch为单位
        pred_choice = pred.data.max(1)[1]
        correct = pred_choice.eq(target.data).cpu().sum()
        print('[%d/%d: %d/%d] train loss: %f accuracy: %f' % (epoch,(opt.nepoch), i, num_batch, loss.item(), correct.item()/float(opt.batchSize * 2500)))

        if i % 10 == 0:  #每10个batch执行一次,即为每10*batchSize个点云执行一次验证
            j, data = next(enumerate(testdataloader, 0))
            points, target = data
            points = points.transpose(2, 1)
            points, target = points.cuda(), target.cuda()
            classifier = classifier.eval()
            pred, _, _ = classifier(points)
            pred = pred.view(-1, num_classes)
            target = target.view(-1, 1)[:, 0] - 1
            loss = F.nll_loss(pred, target)
            # print('pred.shape:{}'.format(pred.size))
            pred_choice = pred.data.max(1)[1]
            correct = pred_choice.eq(target.data).cpu().sum()
            print('[%d/%d: %d/%d] %s loss: %f accuracy: %f' % (epoch,(opt.nepoch), i, num_batch, blue('test'), loss.item(), correct.item()/float(opt.batchSize * 2500)))

    scheduler.step()   #更新一下学习率 每step_size=20调整一次

    torch.save(classifier.state_dict(), '%s/seg_model_%s_%d.pth' % (opt.outf, opt.class_choice, epoch))   #每个epoch都保存一个模型


# 计算精确度
## benchmark mIOU
shape_ious = []
for i,data in tqdm(enumerate(testdataloader, 0)):  # tqdm进度条
    points, target = data
    points = points.transpose(2, 1)
    points, target = points.cuda(), target.cuda()
    classifier = classifier.eval()   #评估
    pred, _, _ = classifier(points)
    print('\npred.shape:{}'.format(pred.shape))  #pred.shape:torch.Size([8, 2500, 4])
    pred_choice = pred.data.max(2)[1]

    pred_np = pred_choice.cpu().data.numpy()
    target_np = target.cpu().data.numpy() - 1
    print('target_np.shape:{}\n'.format(target_np.shape))  #target_np.shape:(8, 2500)


    for shape_idx in range(target_np.shape[0]):
        parts = range(num_classes)#np.unique(target_np[shape_idx])
        part_ious = []
        for part in parts:
            I = np.sum(np.logical_and(pred_np[shape_idx] == part, target_np[shape_idx] == part))
            U = np.sum(np.logical_or(pred_np[shape_idx] == part, target_np[shape_idx] == part))
            if U == 0:
                iou = 1 #If the union of groundtruth and prediction points is empty, then count part IoU as 1
            else:
                iou = I / float(U)
            part_ious.append(iou)
        shape_ious.append(np.mean(part_ious))

print("mIOU for class {}: {}".format(opt.class_choice, np.mean(shape_ious)))

pointnet

├── dataset.py
├── __init__.py
├── model.py
└── __pycache__
    ├── dataset.cpython-36.pyc
    ├── __init__.cpython-36.pyc
    └── model.cpython-36.pyc

1 directory, 6 files

简介

• dataset.py数据集导入相关
• model.py模型构建相关

代码详细注释

• dataset.py

# from __future__ import print_function
import torch.utils.data as data
import os
import os.path
import torch
import numpy as np
import sys
from tqdm import tqdm 
import json
from plyfile import PlyData, PlyElement

def get_segmentation_classes(root):
    catfile = os.path.join(root, 'synsetoffset2category.txt')
    cat = {}
    meta = {}

    with open(catfile, 'r') as f:
        for line in f:
            ls = line.strip().split()
            cat[ls[0]] = ls[1]

    for item in cat:
        dir_seg = os.path.join(root, cat[item], 'points_label')
        dir_point = os.path.join(root, cat[item], 'points')
        fns = sorted(os.listdir(dir_point))
        meta[item] = []
        for fn in fns:
            token = (os.path.splitext(os.path.basename(fn))[0])
            meta[item].append((os.path.join(dir_point, token + '.pts'), os.path.join(dir_seg, token + '.seg')))
    
    with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../misc/num_seg_classes.txt'), 'w') as f:
        for item in cat:
            datapath = []
            num_seg_classes = 0
            for fn in meta[item]:
                datapath.append((item, fn[0], fn[1]))

            for i in tqdm(range(len(datapath))):
                l = len(np.unique(np.loadtxt(datapath[i][-1]).astype(np.uint8)))
                if l > num_seg_classes:
                    num_seg_classes = l

            print("category {} num segmentation classes {}".format(item, num_seg_classes))
            f.write("{}\t{}\n".format(item, num_seg_classes))

def gen_modelnet_id(root):
    classes = []
    with open(os.path.join(root, 'train.txt'), 'r') as f:
        for line in f:
            classes.append(line.strip().split('/')[0])
    classes = np.unique(classes)
    with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../misc/modelnet_id.txt'), 'w') as f:
        for i in range(len(classes)):
            f.write('{}\t{}\n'.format(classes[i], i))


# 导入ShapeNetDataset数据集的类
class ShapeNetDataset(data.Dataset):
    def __init__(self,
                 root,
                 npoints=2500,
                 classification=False,
                 class_choice=None,
                 # split='test',
                 split='train',
                 data_augmentation=True):
        self.npoints = npoints
        self.root = root
        self.catfile = os.path.join(self.root, 'synsetoffset2category.txt')

        # #输出catfile的名称
        # print('\n------------------------\nself.catfile:{}'.format(self.catfile),'\n-------------------------')

        self.cat = {}  #字典
        self.data_augmentation = data_augmentation  #默认1
        self.classification = classification  #默认0
        self.seg_classes = {}   #字典

        #读取synsetoffset2category.txt
        with open(self.catfile, 'r') as f:
            for line in f:
                ls = line.strip().split()  #ls是一个list. ls: ['Airplane', '02691156']
                # print("------------------")
                # print('ls:',ls)
                # print("------------------")
                self.cat[ls[0]] = ls[1]

        #显示这个字典cat{}
        # print('self.cat:',self.cat)  #self.cat: {'Airplane': '02691156', 'Bag': '02773838', 'Cap': '02954340', 'Car': '02958343', 'Chair': '03001627', 'Earphone': '03261776', 'Guitar': '03467517', 'Knife': '03624134', 'Lamp': '03636649', 'Laptop': '03642806', 'Motorbike': '03790512', 'Mug': '03797390', 'Pistol': '03948459', 'Rocket': '04099429', 'Skateboard': '04225987', 'Table': '04379243'}

        # print('self.cat.items() :', self.cat.items())   #cat.items()返回一个list[],元素为元组().即[(),(),()....]
        if not class_choice is None:  #当class_choice不是None的时候才会执行
            self.cat = {k: v for k, v in self.cat.items() if k in class_choice}  #self.cat:{'Airplane': '02691156'}
        # print('-------------------------')
        # print('self.cat:{}'.format(self.cat))
        # print('-------------------------')

        self.id2cat = {v: k for k, v in self.cat.items()}
        # print('self.id2cat:{}'.format(self.id2cat))  #self.id2cat:{'02691156': 'Airplane'}

        self.meta = {} #字典
        splitfile = os.path.join(self.root, 'train_test_split', 'shuffled_{}_file_list.json'.format(split))   #读取shuffled_train_file_list.json文件
        # print('splitfile:{}'.format(splitfile))

        #from IPython import embed; embed()
        filelist = json.load(open(splitfile, 'r')) #读取shuffled_train_file_list.json文件
        for item in self.cat:  #self.cat:{'Airplane': '02691156'}
            # print('item:{}'.format(item))  #item:Airplane
            self.meta[item] = []
        # print('self.meta:{}'.format(self.meta))  #self.meta:{'Airplane': []}

        # print('self.cat.values():{}'.format(self.cat.values()))
        for file in filelist:
            # print('file:{}'.format(file))
            _, category, uuid = file.split('/')  #获取每一条的shape_data/04379243/9e3f1901ea14aca753315facdf531a34

            if category in self.cat.values():
                # print("category:{}".format(category))   #04379243

                self.meta[self.id2cat[category]].append((os.path.join(self.root, category, 'points', uuid+'.pts'),
                                        os.path.join(self.root, category, 'points_label', uuid+'.seg')))
        # print("self.meta:{}".format(self.meta))   #{'Airplane':[(),(),()....() ] }
        self.datapath = []  #list
        for item in self.cat:  #self.cat:{'Airplane': '02691156'}
            # print('item:{}'.format(item))  #item:Airplane
            for fn in self.meta[item]:
                self.datapath.append((item, fn[0], fn[1]))   #datapath:[(),()...],其中()为('Airplane','.pts','.seg')元组

        # print('----------------------')
        # print("datapath:{}".format(self.datapath))
        # print('----------------((------')
        # print('cat:{}'.format(self.cat))  #cat:{'Airplane': '02691156'}

        # print('len(self.cat):{}'.format(len(self.cat)))  #len(self.cat):1
        self.classes = dict(zip(sorted(self.cat), range(len(self.cat))))  #使用zip函数, 把key和value的list组合在一起, 再转成字典(dict).
        # print(self.classes) #{'Airplane': 0}
        with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../misc/num_seg_classes.txt'), 'r') as f:
            for line in f:
                ls = line.strip().split()   #ls是一个list[]
                self.seg_classes[ls[0]] = int(ls[1]) #seg_classes是一个字典{}  ->seg_classes:{'Airplane': 4, 'Bag': 2, 'Cap': 2, 'Car': 4, 'Chair': 4, 'Earphone': 3, 'Guitar': 3, 'Knife': 2, 'Lamp': 4, 'Laptop': 2, 'Motorbike': 6, 'Mug': 2, 'Pistol': 3, 'Rocket': 3, 'Skateboard': 3, 'Table': 3}

        # print('seg_classes:{}'.format(self.seg_classes))
        # print('self.cat.keys():{}'.format(self.cat.keys()))  #self.cat.keys():dict_keys(['Airplane'])
        # print('list(self.cat.keys())[0]:{}'.format(list(self.cat.keys())[0]))  #转化为一个list->list(self.cat.keys()):['Airplane']
        self.num_seg_classes = self.seg_classes[list(self.cat.keys())[0]]   #num_seg_classes为对应的的类应该分成几类

        # 输出:{'Airplane': 4, 'Bag': 2, 'Cap': 2, 'Car': 4, 'Chair': 4, 'Earphone': 3, 'Guitar': 3, 'Knife': 2, 'Lamp': 4, 'Laptop': 2, 'Motorbike': 6, 'Mug': 2, 'Pistol': 3, 'Rocket': 3, 'Skateboard': 3, 'Table': 3} 4
        print(self.seg_classes, list(self.cat.keys())[0],self.num_seg_classes)  #输出:一个字典{} 一个数字

    def __getitem__(self, index):
        fn = self.datapath[index]  #fn为元组()  ()为('Airplane','.pts','.seg')元组
        cls = self.classes[self.datapath[index][0]] #cls为classes{}的value
        # print('cls:{}'.format(cls))

        #读取点云和分类标签
        point_set = np.loadtxt(fn[1]).astype(np.float32)
        seg = np.loadtxt(fn[2]).astype(np.int64)
        # print('point_set.shape:{}_____ seg.shape:{}'.format(point_set.shape, seg.shape))  #point_set.shape:(2658, 3)_____ seg.shape:(2658,)

# 重新采样到self.npoints个点
        choice = np.random.choice(len(seg), self.npoints, replace=True)
        #resample
        point_set = point_set[choice, :]

        # print('point_set:{}'.format(point_set))
        # print('np.mean(point_set, axis = 0):{}'.format(np.mean(point_set, axis = 0)))  #中心点 1*3
        point_set = point_set - np.expand_dims(np.mean(point_set, axis = 0), 0) # 去中心化
        # print('new_point_set:{}'.format(point_set))

        dist = np.max(np.sqrt(np.sum(point_set ** 2, axis = 1)),0)  #计算到原点的最远距离
        # print('dist:{}'.format(dist))
        point_set = point_set / dist #scale  归一化

        if self.data_augmentation:  #默认False  开启旋转任意角度并加上一个bias
            theta = np.random.uniform(0,np.pi*2)
            rotation_matrix = np.array([[np.cos(theta), -np.sin(theta)],[np.sin(theta), np.cos(theta)]])
            # print('rotation_matrix:{}'.format(rotation_matrix))
            point_set[:,[0,2]] = point_set[:,[0,2]].dot(rotation_matrix) # random rotation
            # print('point_set:{}'.format(point_set.shape))
            point_set += np.random.normal(0, 0.02, size=point_set.shape) # random jitter

        seg = seg[choice]
        point_set = torch.from_numpy(point_set)
        seg = torch.from_numpy(seg)
        cls = torch.from_numpy(np.array([cls]).astype(np.int64))  #cls为对应的代号,比如Airplane对应0



        if self.classification:  #classification默认是0
            return point_set, cls
        else:
            return point_set, seg

    def __len__(self):
        # print('len(self.datapath):{}'.format(len(self.datapath)))
        return len(self.datapath)

class ModelNetDataset(data.Dataset):
    def __init__(self,
                 root,
                 npoints=2500,
                 split='train',
                 data_augmentation=True):
        self.npoints = npoints
        self.root = root
        self.split = split
        self.data_augmentation = data_augmentation
        self.fns = []
        with open(os.path.join(root, '{}.txt'.format(self.split)), 'r') as f:
            for line in f:
                self.fns.append(line.strip())

        self.cat = {}
        with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../misc/modelnet_id.txt'), 'r') as f:
            for line in f:
                ls = line.strip().split()
                self.cat[ls[0]] = int(ls[1])

        print(self.cat)
        self.classes = list(self.cat.keys())

    def __getitem__(self, index):
        fn = self.fns[index]
        cls = self.cat[fn.split('/')[0]]
        with open(os.path.join(self.root, fn), 'rb') as f:
            plydata = PlyData.read(f)
        pts = np.vstack([plydata['vertex']['x'], plydata['vertex']['y'], plydata['vertex']['z']]).T
        choice = np.random.choice(len(pts), self.npoints, replace=True)
        point_set = pts[choice, :]

        point_set = point_set - np.expand_dims(np.mean(point_set, axis=0), 0)  # center
        dist = np.max(np.sqrt(np.sum(point_set ** 2, axis=1)), 0)
        point_set = point_set / dist  # scale

        if self.data_augmentation:
            theta = np.random.uniform(0, np.pi * 2)
            rotation_matrix = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]])
            point_set[:, [0, 2]] = point_set[:, [0, 2]].dot(rotation_matrix)  # random rotation
            point_set += np.random.normal(0, 0.02, size=point_set.shape)  # random jitter

        point_set = torch.from_numpy(point_set.astype(np.float32))
        cls = torch.from_numpy(np.array([cls]).astype(np.int64))
        return point_set, cls


    def __len__(self):
        return len(self.fns)

if __name__ == '__main__':
    dataset = sys.argv[1]
    datapath = sys.argv[2]

    if dataset == 'shapenet':
        d = ShapeNetDataset(root = datapath, class_choice = ['Bag'])
        print('len(d):{}'.format(len(d)))
        ps, seg = d[0]
        print(ps.size(), ps.type(), seg.size(),seg.type())

        print('--------------------------------------------------------')
        d = ShapeNetDataset(root = datapath, classification = True)
        print('len(d):{}'.format(len(d)))
        ps, cls = d[0]
        print(ps.size(), ps.type(), cls.size(),cls.type())
        # get_segmentation_classes(datapath)

    if dataset == 'modelnet':
        gen_modelnet_id(datapath)
        d = ModelNetDataset(root=datapath)
        print(len(d))
        print(d[0])

• model.py

• 包含以下结构

  • 模型
    • STN3d
    • STNkd
    • PointNetfeat
    • PointNetCls ---- 分类网络
    • PointNetDenseCls ------分割网络
  • 函数
    • feature_transform_regularizer

# from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.utils.data
from torch.autograd import Variable
import numpy as np
import torch.nn.functional as F


class STN3d(nn.Module):
    def __init__(self):
        super(STN3d, self).__init__()
        self.conv1 = torch.nn.Conv1d(3, 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, 9)
        self.relu = nn.ReLU()

        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, x):
        batchsize = x.size()[0]
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = torch.max(x, 2, keepdim=True)[0]
        x = x.view(-1, 1024)

        x = F.relu(self.bn4(self.fc1(x)))
        x = F.relu(self.bn5(self.fc2(x)))
        x = self.fc3(x)

        iden = Variable(torch.from_numpy(np.array([1,0,0,0,1,0,0,0,1]).astype(np.float32))).view(1,9).repeat(batchsize,1)
        if x.is_cuda:
            iden = iden.cuda()
        x = x + iden
        x = x.view(-1, 3, 3)
        return x


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.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(1024)
        self.bn4 = nn.BatchNorm1d(512)
        self.bn5 = nn.BatchNorm1d(256)

        self.k = k

    def forward(self, x):
        batchsize = x.size()[0]
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = torch.max(x, 2, keepdim=True)[0]
        x = x.view(-1, 1024)

        x = F.relu(self.bn4(self.fc1(x)))
        x = F.relu(self.bn5(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.cuda()
        x = x + iden
        x = x.view(-1, self.k, self.k)
        return x

class PointNetfeat(nn.Module):
    def __init__(self, global_feat = True, feature_transform = False):
        super(PointNetfeat, self).__init__()
        self.stn = STN3d()
        self.conv1 = torch.nn.Conv1d(3, 64, 1)
        self.conv2 = torch.nn.Conv1d(64, 128, 1)
        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(1024)
        self.global_feat = global_feat
        self.feature_transform = feature_transform
        if self.feature_transform:
            self.fstn = STNkd(k=64)

    def forward(self, x):
        n_pts = x.size()[2]
        trans = self.stn(x)
        x = x.transpose(2, 1)
        x = torch.bmm(x, trans)
        x = x.transpose(2, 1)
        x = F.relu(self.bn1(self.conv1(x)))

        if self.feature_transform:
            trans_feat = self.fstn(x)
            x = x.transpose(2,1)
            x = torch.bmm(x, trans_feat)
            x = x.transpose(2,1)
        else:
            trans_feat = None

        pointfeat = x
        x = F.relu(self.bn2(self.conv2(x)))
        x = self.bn3(self.conv3(x))
        x = torch.max(x, 2, keepdim=True)[0]
        x = x.view(-1, 1024)
        if self.global_feat:
            return x, trans, trans_feat
        else:
            x = x.view(-1, 1024, 1).repeat(1, 1, n_pts)
            return torch.cat([x, pointfeat], 1), trans, trans_feat

class PointNetCls(nn.Module):
    def __init__(self, k=2, feature_transform=False):
        super(PointNetCls, self).__init__()
        self.feature_transform = feature_transform
        self.feat = PointNetfeat(global_feat=True, feature_transform=feature_transform)
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, k)
        self.dropout = nn.Dropout(p=0.3)
        self.bn1 = nn.BatchNorm1d(512)
        self.bn2 = nn.BatchNorm1d(256)
        self.relu = nn.ReLU()

    def forward(self, x):
        x, trans, trans_feat = self.feat(x)
        x = F.relu(self.bn1(self.fc1(x)))
        x = F.relu(self.bn2(self.dropout(self.fc2(x))))
        x = self.fc3(x)
        return F.log_softmax(x, dim=1), trans, trans_feat


class PointNetDenseCls(nn.Module):
    def __init__(self, k = 2, feature_transform=False):
        super(PointNetDenseCls, self).__init__()
        self.k = k
        self.feature_transform=feature_transform
        self.feat = PointNetfeat(global_feat=False, feature_transform=feature_transform)
        self.conv1 = torch.nn.Conv1d(1088, 512, 1)
        self.conv2 = torch.nn.Conv1d(512, 256, 1)
        self.conv3 = torch.nn.Conv1d(256, 128, 1)
        self.conv4 = torch.nn.Conv1d(128, self.k, 1)
        self.bn1 = nn.BatchNorm1d(512)
        self.bn2 = nn.BatchNorm1d(256)
        self.bn3 = nn.BatchNorm1d(128)

    def forward(self, x):
        batchsize = x.size()[0]
        n_pts = x.size()[2]
        x, trans, trans_feat = self.feat(x)
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = self.conv4(x)

        # print('x:{}'.format(x.shape))  #x:torch.Size([32, 3, 2500])
        x = x.transpose(1,2).contiguous()   #transpose()交换1和2两个维度    contiguous()用来将tensor变为连续的
        # print('x.transpose(2,1):{}'.format(x.shape))  #x.transpose(2,1):torch.Size([32, 2500, 3])


        x = F.log_softmax(x.view(-1,self.k), dim=-1)
        # print('log_softmax:{}'.format(x.shape))  #torch.Size([20000, 4])
        x = x.view(batchsize, n_pts, self.k)
        # print('x.shape:{}'.format(x.shape))
        return x, trans, trans_feat

def feature_transform_regularizer(trans):
    d = trans.size()[1]
    batchsize = trans.size()[0]
    I = torch.eye(d)[None, :, :]   # eye:产生一个单位矩阵
    if trans.is_cuda:
        I = I.cuda()
    loss = torch.mean(torch.norm(torch.bmm(trans, trans.transpose(2,1)) - I, dim=(1,2)))
    return loss

if __name__ == '__main__':
    sim_data = Variable(torch.rand(32,3,2500))
    trans = STN3d()
    out = trans(sim_data)
    print('stn', out.size())
    print('loss', feature_transform_regularizer(out))

    sim_data_64d = Variable(torch.rand(32, 64, 2500))
    trans = STNkd(k=64)
    out = trans(sim_data_64d)
    print('stn64d', out.size())
    print('loss', feature_transform_regularizer(out))

    pointfeat = PointNetfeat(global_feat=True)
    out, _, _ = pointfeat(sim_data)
    print('global feat', out.size())

    pointfeat = PointNetfeat(global_feat=False)
    out, _, _ = pointfeat(sim_data)
    print('point feat', out.size())

    cls = PointNetCls(k = 5)
    out, _, _ = cls(sim_data)
    print('class', out.size())

    seg = PointNetDenseCls(k = 3)
    out, _, _ = seg(sim_data)
    print('seg', out.size())

    print(seg)

Reference

• 言简意赅python系列—if not x: 和 if x is not None: 和 if not x is None: 的区别
• Python 字典(Dictionary) items()方法
• Python - 两个列表(list)组成字典(dict)
• Python.__getitem__方法
• python四个魔法方法__len__,getitem,setitem,delitem
  

开源框架PointNet 代码详解——/pointnet/sem_seg/train.py
【3D计算机视觉】从PointNet到PointNet++理论及pytorch代码

Tools

• 每次5秒更新一次gpu状态

watch -n 5 nvidia-smi