三维点云学习（5）5-实现Deeplearning-PointNet-2-classfication
Github PointNet源码
数据集下载：为40种物体的三维点云数据集
提取码：es14
运行
Epoch=50后的结果
trainclassification.py
#train_classification.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
#from pointnet-pytorch-master.pointnet.dataset import ShapeNetDataset, ModelNetDataset
#from pointnet-pytorch-master.pointnet.model import PointNetCls, feature_transform_regularizer
import torch.nn.functional as F
from tqdm import tqdm
import sys;
sys.path.append('/home/renzhanqi/workspace/studyMaterialsAndNotes/shenLanLidarProcess/shenLanLidarPrcess-优快云-Materials/code/pointnet-pytorch-master/pointnet');
from dataset_modify import ShapeNetDataset, ModelNetDataset;
from model import PointNetCls, feature_transform_regularizer
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument(
'--batchSize', type=int, default=32, help='input batch size')
parser.add_argument(
'--num_points', type=int, default=2500, help='input batch size')
parser.add_argument(
'--workers', type=int, help='number of data loading workers', default=4)
parser.add_argument(
'--nepoch', type=int, default=250, help='number of epochs to train for')
parser.add_argument('--outf', type=str, default='cls', 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('--dataset_type', type=str, default='shapenet', help="dataset type shapenet|modelnet40")
parser.add_argument('--feature_transform', action='store_true', help="use feature transform")
opt = parser.parse_args()
print(opt)
blue = lambda x: '\033[94m' + x + '\033[0m'
opt.manualSeed = random.randint(1, 10000) # fix seed
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if opt.dataset_type == 'shapenet':
dataset = ShapeNetDataset(
root=opt.dataset,
classification=True,
npoints=opt.num_points)
test_dataset = ShapeNetDataset(
root=opt.dataset,
classification=True,
split='test',
npoints=opt.num_points,
data_augmentation=False)
elif opt.dataset_type == 'modelnet40':
dataset = ModelNetDataset(
root=opt.dataset,
npoints=opt.num_points,
split='train')
test_dataset = ModelNetDataset(
root=opt.dataset,
split='test',
npoints=opt.num_points,
data_augmentation=False)
else:
exit('wrong dataset type')
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=opt.batchSize,
shuffle=True,
num_workers=int(opt.workers))
testdataloader = torch.utils.data.DataLoader(
test_dataset,
batch_size=opt.batchSize,
shuffle=True,
num_workers=int(opt.workers))
print(len(dataset), len(test_dataset))
num_classes = len(dataset.classes)
print('classes', num_classes)
try:
os.makedirs(opt.outf)
except OSError:
pass
classifier = PointNetCls(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=20, gamma=0.5)
classifier.cuda()
num_batch = len(dataset) / opt.batchSize
for epoch in range(opt.nepoch):
scheduler.step()
for i, data in enumerate(dataloader, 0):
points, target = data
target = target[:, 0]
points = points.transpose(2, 1)
points, target = points.cuda(), target.cuda()
optimizer.zero_grad()
classifier = classifier.train()
pred, trans, trans_feat = classifier(points)
loss = F.nll_loss(pred, target)
if opt.feature_transform:
loss += feature_transform_regularizer(trans_feat) * 0.001
loss.backward()
optimizer.step()
pred_choice = pred.data.max(1)[1]
correct = pred_choice.eq(target.data).cpu().sum()
print('[%d: %d/%d] train loss: %f accuracy: %f' % (epoch, i, num_batch, loss.item(), correct.item() / float(opt.batchSize)))
if i % 10 == 0:
j, data = next(enumerate(testdataloader, 0))
points, target = data
target = target[:, 0]
points = points.transpose(2, 1)
points, target = points.cuda(), target.cuda()
classifier = classifier.eval()
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('[%d: %d/%d] %s loss: %f accuracy: %f' % (epoch, i, num_batch, blue('test'), loss.item(), correct.item()/float(opt.batchSize)))
torch.save(classifier.state_dict(), '%s/cls_model_%d.pth' % (opt.outf, epoch))
total_correct = 0
total_testset = 0
for i,data in tqdm(enumerate(testdataloader, 0)):
points, target = data
target = target[:, 0]
points = points.transpose(2, 1)
points, target = points.cuda(), target.cuda()
classifier = classifier.eval()
pred, _, _ = classifier(points)
pred_choice = pred.data.max(1)[1]
correct = pred_choice.eq(target.data).cpu().sum()
total_correct += correct.item()
total_testset += points.size()[0]
print("final accuracy {}".format(total_correct / float(total_testset)))
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))
class ShapeNetDataset(data.Dataset):
def __init__(self,
root,
npoints=2500,
classification=False,
class_choice=None,
split='train',
data_augmentation=True):
self.npoints = npoints
self.root = root
self.catfile = os.path.join(self.root, 'synsetoffset2category.txt')
self.cat = {}
self.data_augmentation = data_augmentation
self.classification = classification
self.seg_classes = {}
with open(self.catfile, 'r') as f:
for line in f:
ls = line.strip().split()
self.cat[ls[0]] = ls[1]
#print(self.cat)
if not class_choice is None:
self.cat = {k: v for k, v in self.cat.items() if k in class_choice}
self.id2cat = {v: k for k, v in self.cat.items()}
self.meta = {}
splitfile = os.path.join(self.root, 'train_test_split', 'shuffled_{}_file_list.json'.format(split))
#from IPython import embed; embed()
filelist = json.load(open(splitfile, 'r'))
for item in self.cat:
self.meta[item] = []
for file in filelist:
_, category, uuid = file.split('/')
if category in self.cat.values():
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')))
self.datapath = []
for item in self.cat:
for fn in self.meta[item]:
self.datapath.append((item, fn[0], fn[1]))
self.classes = dict(zip(sorted(self.cat), range(len(self.cat))))
print(self.classes)
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()
self.seg_classes[ls[0]] = int(ls[1])
self.num_seg_classes = self.seg_classes[list(self.cat.keys())[0]]
print(self.seg_classes, self.num_seg_classes)
def __getitem__(self, index):
fn = self.datapath[index]
cls = self.classes[self.datapath[index][0]]
point_set = np.loadtxt(fn[1]).astype(np.float32)
seg = np.loadtxt(fn[2]).astype(np.int64)
#print(point_set.shape, seg.shape)
choice = np.random.choice(len(seg), self.npoints, replace=True)
#resample
point_set = point_set[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
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))
if self.classification:
return point_set, cls
else:
return point_set, seg
def __len__(self):
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.files_names = []
with open(os.path.join(root, "modelnet40_{}.txt".format(self.split))) as file:
for line in file:
self.files_names.append(line.strip()) # e.g. "airplane_0001"
# 读取类别映射
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):
file_index = self.files_names[index] # eg: index:845 -> bed_0114
cls_name = file_index.rsplit(sep='_', maxsplit=1)[0] # eg #airplane，rsplit:右侧分割，maxplit：限定分割符的出现次数
cls = self.cat[cls_name] # eg out:[0] airplane属于[0],bench属于[3]
file_name = '{}/{}.txt'.format(cls_name, file_index) # eg airplane/airplane_0001.txt
pts = np.loadtxt(os.path.join(self.root, file_name), delimiter=',', dtype=float)[:,
:3] # 读取每个point的前三个数字，分别为(x,y,z);后三位为法向量不读取
# 随机降采样，提高模型的鲁棒性
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.files_names)
if __name__ == '__main__':
dataset = sys.argv[1]
datapath = sys.argv[2]
if dataset == 'shapenet':
d = ShapeNetDataset(root = datapath, class_choice = ['Chair'])
print(len(d))
ps, seg = d[0]
print(ps.size(), ps.type(), seg.size(),seg.type())
d = ShapeNetDataset(root = datapath, classification = True)
print(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])
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
