跑自己的点云数据——详细代码与注解
应之前小伙伴的要求，这里给出详细的修改后的点云分割代码，可以跑通自己的点云数据。
首先给出训练部分的代码：（这里用到的是：train——partseg）
import argparseimport osfrom data_utils.ShapeNetDataLoader import PartNormalDatasetimport torchimport datetimeimport loggingfrom pathlib import Pathimport sysimport importlibimport shutilfrom tqdm import tqdmimport providerimport numpy as np"""训练所需设置参数：--model pointnet2_part_seg_msg --normal --log_dir pointnet2_part_seg_msg"""BASE_DIR = os.path.dirname(os.path.abspath(__file__))ROOT_DIR = BASE_DIRsys.path.append(os.path.join(ROOT_DIR, 'models'))##seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43], 'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46], 'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27], 'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40], 'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}seg_classes = {'Airplane': [0, 1], 'Mug': [2, 3]}seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table}for cat in seg_classes.keys(): for label in seg_classes[cat]: seg_label_to_cat[label] = catdef to_categorical(y, num_classes): """ 1-hot encodes a tensor """ new_y = torch.eye(num_classes)[y.cpu().data.numpy(),] if (y.is_cuda): return new_y.cuda() return new_ydef parse_args(): parser = argparse.ArgumentParser('Model') parser.add_argument('--model', type=str, default='pointnet2_part_seg_msg', help='model name [default: pointnet2_part_seg_msg]') parser.add_argument('--batch_size', type=int, default=4, help='Batch Size during training [default: 16]') parser.add_argument('--epoch', default=251, type=int, help='Epoch to run [default: 251]') parser.add_argument('--learning_rate', default=0.001, type=float, help='Initial learning rate [default: 0.001]') parser.add_argument('--gpu', type=str, default='0', help='GPU to use [default: GPU 0]') parser.add_argument('--optimizer', type=str, default='Adam', help='Adam or SGD [default: Adam]') parser.add_argument('--log_dir', type=str, default=None, help='Log path [default: None]') parser.add_argument('--decay_rate', type=float, default=1e-4, help='weight decay [default: 1e-4]') parser.add_argument('--npoint', type=int, default=2048, help='Point Number [default: 2048]') parser.add_argument('--normal', action='store_true', default=False, help='Whether to use normal information [default: False]') parser.add_argument('--step_size', type=int, default=2, help='Decay step for lr decay [default: every 20 epochs]') parser.add_argument('--lr_decay', type=float, default=0.5, help='Decay rate for lr decay [default: 0.5]') return parser.parse_args()def main(args): def log_string(str): logger.info(str) print(str) # '''HYPER PARAMETER''' # os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) experiment_dir = Path('./log/') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath('part_seg') experiment_dir.mkdir(exist_ok=True) if args.log_dir is None: experiment_dir = experiment_dir.joinpath(timestr) else: experiment_dir = experiment_dir.joinpath(args.log_dir) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = experiment_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger("Model") logger.setLevel(logging.INFO) formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) log_string('PARAMETER ...') log_string(args) root = 'data/shapenetcore_partanno_segmentation_benchmark_v0_normal/' TRAIN_DATASET = PartNormalDataset(root = root, npoints=args.npoint, split='trainval', normal_channel=args.normal) trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size,shuffle=True, num_workers=4) TEST_DATASET = PartNormalDataset(root = root, npoints=args.npoint, split='test', normal_channel=args.normal) testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size,shuffle=False, num_workers=4) log_string("The number of training data is: %d" % len(TRAIN_DATASET)) log_string("The number of test data is: %d" % len(TEST_DATASET)) num_classes = 2 num_part = 4 '''MODEL LOADING''' MODEL = importlib.import_module(args.model) shutil.copy('models/%s.py' % args.model, str(experiment_dir)) shutil.copy('models/pointnet_util.py', str(experiment_dir)) classifier = MODEL.get_model(num_part, normal_channel=args.normal).cuda() criterion = MODEL.get_loss().cuda() def weights_init(m): classname = m.__class__.__name__ if classname.find('Conv2d') != -1: torch.nn.init.xavier_normal_(m.weight.data) torch.nn.init.constant_(m.bias.data, 0.0) elif classname.find('Linear') != -1: torch.nn.init.xavier_normal_(m.weight.data) torch.nn.init.constant_(m.bias.data, 0.0) try: checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth') start_epoch = checkpoint['epoch'] classifier.load_state_dict(checkpoint['model_state_dict']) log_string('Use pretrain model') except: log_string('No existing model, starting training from scratch...') start_epoch = 0 classifier = classifier.apply(weights_init) if args.optimizer == 'Adam': optimizer = torch.optim.Adam( classifier.parameters(), lr=args.learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=args.decay_rate ) else: optimizer = torch.optim.SGD(classifier.parameters(), lr=args.learning_rate, momentum=0.9) def bn_momentum_adjust(m, momentum): if isinstance(m, torch.nn.BatchNorm2d) or isinstance(m, torch.nn.BatchNorm1d): m.momentum = momentum LEARNING_RATE_CLIP = 1e-5 MOMENTUM_ORIGINAL = 0.1 MOMENTUM_DECCAY = 0.5 MOMENTUM_DECCAY_STEP = args.step_size best_acc = 0 global_epoch = 0 best_class_avg_iou = 0 best_inctance_avg_iou = 0 for epoch in range(start_epoch,args.epoch): log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch)) '''Adjust learning rate and BN momentum''' lr = max(args.learning_rate * (args.lr_decay ** (epoch // args.step_size)), LEARNING_RATE_CLIP) log_string('Learning rate:%f' % lr) for param_group in optimizer.param_groups: param_group['lr'] = lr mean_correct = [] momentum = MOMENTUM_ORIGINAL * (MOMENTUM_DECCAY ** (epoch // MOMENTUM_DECCAY_STEP)) if momentum = best_inctance_avg_iou): logger.info('Save model...') savepath = str(checkpoints_dir) + '/best_model.pth' log_string('Saving at %s'% savepath) state = { 'epoch': epoch, 'train_acc': train_instance_acc, 'test_acc': test_metrics['accuracy'], 'class_avg_iou': test_metrics['class_avg_iou'], 'inctance_avg_iou': test_metrics['inctance_avg_iou'], 'model_state_dict': classifier.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), } torch.save(state, savepath) log_string('Saving model....') if test_metrics['accuracy'] > best_acc: best_acc = test_metrics['accuracy'] if test_metrics['class_avg_iou'] > best_class_avg_iou: best_class_avg_iou = test_metrics['class_avg_iou'] if test_metrics['inctance_avg_iou'] > best_inctance_avg_iou: best_inctance_avg_iou = test_metrics['inctance_avg_iou'] log_string('Best accuracy is: %.5f'%best_acc) log_string('Best class avg mIOU is: %.5f'%best_class_avg_iou) log_string('Best inctance avg mIOU is: %.5f'%best_inctance_avg_iou) global_epoch+=1if __name__ == '__main__': args = parse_args() main(args)接着是测试部分的代码：（用到的是：test——partseg）
import argparseimport osfrom data_utils.ShapeNetDataLoader import PartNormalDatasetimport torchimport loggingimport sysimport importlibfrom tqdm import tqdmimport numpy as np#我的电脑不支持cuda，所以我的文件把所有cuda删掉了，这个文件又把他们都加回来了。BASE_DIR = os.path.dirname(os.path.abspath(__file__))ROOT_DIR = BASE_DIRsys.path.append(os.path.join(ROOT_DIR, 'models'))seg_classes = {'Airplane': [0, 1], 'Mug': [2, 3]}#我们这里是将两个类四个部件分割seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table}for cat in seg_classes.keys(): for label in seg_classes[cat]: seg_label_to_cat[label] = catdef to_categorical(y, num_classes): """ 1-hot encodes a tensor """ new_y = torch.eye(num_classes)[y.cpu().data.numpy(),] if (y.is_cuda): return new_y.cuda() return new_ydef parse_args(): '''PARAMETERS''' parser = argparse.ArgumentParser('PointNet') parser.add_argument('--batch_size', type=int, default=250, help='batch size in testing [default: 24]')#这里由于输出循环没太难弄懂，所以直接236（250大于236，设置一个比236大的就可以了）个测试集弄在一个batch中跑，我试过了可以跑，慢一点而已，但是如果test的太多了就估计不行了 parser.add_argument('--gpu', type=str, default='0', help='specify gpu device [default: 0]') parser.add_argument('--num_point', type=int, default=2048, help='Point Number [default: 2048]') parser.add_argument('--log_dir', type=str, default='pointnet2_part_seg_ssg', help='Experiment root') parser.add_argument('--normal', action='store_true', default=False, help='Whether to use normal information [default: False]') parser.add_argument('--num_votes', type=int, default=3, help='Aggregate segmentation scores with voting [default: 3]') return parser.parse_args()def main(args): xxxxxx = 0#看一共有多少个点的，没啥用 def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu experiment_dir = 'log/part_seg/' + args.log_dir '''LOG''' args = parse_args() logger = logging.getLogger("Model") logger.setLevel(logging.INFO) formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/eval.txt' % experiment_dir) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) log_string('PARAMETER ...') log_string(args) root = 'data/shapenetcore_partanno_segmentation_benchmark_v0_normal/' TEST_DATASET = PartNormalDataset(root = root, npoints=args.num_point, split='test', normal_channel=args.normal) testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size,shuffle=False, num_workers=4) log_string("The number of test data is: %d" % len(TEST_DATASET)) num_classes = 2 #这里2，4要根据具体的情况来改 num_part = 4 '''MODEL LOADING''' model_name = os.listdir(experiment_dir+'/logs')[0].split('.')[0] MODEL = importlib.import_module(model_name) classifier = MODEL.get_model(num_part, normal_channel=args.normal).cuda() checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth') classifier.load_state_dict(checkpoint['model_state_dict']) with torch.no_grad(): test_metrics = {} total_correct = 0 total_seen = 0 total_seen_class = [0 for _ in range(num_part)] total_correct_class = [0 for _ in range(num_part)] shape_ious = {cat: [] for cat in seg_classes.keys()} seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table} for cat in seg_classes.keys(): for label in seg_classes[cat]: seg_label_to_cat[label] = cat for batch_id, (points, label, target) in tqdm(enumerate(testDataLoader), total=len(testDataLoader), smoothing=0.9): batchsize, num_point, _ = points.size() cur_batch_size, NUM_POINT, _ = points.size() points, label, target = points.float().cuda(), label.long().cuda(), target.long().cuda() points = points.transpose(2, 1) classifier = classifier.eval() vote_pool = torch.zeros(target.size()[0], target.size()[1], num_part).cuda() for _ in range(args.num_votes): seg_pred, _ = classifier(points, to_categorical(label, num_classes)) vote_pool += seg_pred seg_pred = vote_pool / args.num_votes cur_pred_val = seg_pred.cpu().data.numpy() cur_pred_val_logits = cur_pred_val cur_pred_val = np.zeros((cur_batch_size, NUM_POINT)).astype(np.int32) target = target.cpu().data.numpy() points1 = points.transpose(2, 1).numpy() for i in range(cur_batch_size): cat = seg_label_to_cat[target[i, 0]] logits = cur_pred_val_logits[i, :, :] cur_pred_val[i, :] = np.argmax(logits[:, seg_classes[cat]], 1) + seg_classes[cat][0] #从这里开始就是依次去打印分类出的四类点的坐标了，下面的txt会自动生成，不用自己去创建空的txt，并且每一类是啥已经可以由txt的名字看出 #这里需要改的就只有下面四个存取文件的路径，我是都改在一个文件下下的。 #我采用的打开文件的方式是a就是往后添加，必须是这个，因为这个循环里必须一点一点添加才能获取所有的，改成w+就只能得到一个点，就是最后一次更新时候的点。 #所以如果在运行完程序后一次想要再运行时候，就要把该目录下的所有文件删掉在运行，不然就换目录。 aaa = numpy.argwhere(cur_pred_val[i] == 0) for j in aaa: # print(points1[i,j]) res1 = open(r'E:\02691156_0_' + str(i) + '.txt', 'a')#02691156_0_xxx表示02691156这个东西的中的标签为0的那些点，xxx就表示在json文件中他是第几个，因为是每一个文件都要单独把0，1分开 res1.write('\n' + str(points1[i, j])) res1.close() xxxxxx = xxxxxx + 1 bbb = numpy.argwhere(cur_pred_val[i] == 1) for j in bbb: # print(points1[i, j]) res2 = open(r'E:\02691156_1_' + str(i) + '.txt', 'a')#同理，这里是02691156这个东西的标签为1的那些点 res2.write('\n' + str(points1[i, j])) res2.close() xxxxxx = xxxxxx + 1#测试的时候看一共有多少个点的，没啥用 ccc = numpy.argwhere(cur_pred_val[i] == 2) for j in ccc: # print(points1[i, j]) res3 = open(r'E:\03797390_2_' + str(i) + '.txt', 'a')#这里是03797390中标签为2的那些点 res3.write('\n' + str(points1[i, j])) res3.close() xxxxxx = xxxxxx + 1 ddd = numpy.argwhere(cur_pred_val[i] == 3) for j in ddd: # print(points1[i, j]) res4 = open(r'E:\03797390_3_' + str(i) + '.txt', 'a')#这里是03797390中标签为3的那些点 res4.write('\n' + str(points1[i, j])) res4.close() xxxxxx = xxxxxx + 1 print("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx")#稍微分割一下，测试的时候看的，没啥用 correct = np.sum(cur_pred_val == target) total_correct += correct total_seen += (cur_batch_size * NUM_POINT) for l in range(num_part): total_seen_class[l] += np.sum(target == l) total_correct_class[l] += (np.sum((cur_pred_val == l) & (target == l))) for i in range(cur_batch_size): segp = cur_pred_val[i, :] segl = target[i, :] cat = seg_label_to_cat[segl[0]] part_ious = [0.0 for _ in range(len(seg_classes[cat]))] for l in seg_classes[cat]: if (np.sum(segl == l) == 0) and (np.sum(segp == l) == 0): # part is not present, no prediction as well part_ious[l - seg_classes[cat][0]] = 1.0 else: part_ious[l - seg_classes[cat][0]] = np.sum((segl == l) & (segp == l)) / float(np.sum((segl == l) | (segp == l))) shape_ious[cat].append(np.mean(part_ious)) all_shape_ious = [] for cat in shape_ious.keys(): for iou in shape_ious[cat]: all_shape_ious.append(iou) shape_ious[cat] = np.mean(shape_ious[cat]) mean_shape_ious = np.mean(list(shape_ious.values())) test_metrics['accuracy'] = total_correct / float(total_seen) test_metrics['class_avg_accuracy'] = np.mean( np.array(total_correct_class) / np.array(total_seen_class, dtype=np.float)) for cat in sorted(shape_ious.keys()): log_string('eval mIoU of %s %f' % (cat + ' ' * (14 - len(cat)), shape_ious[cat])) test_metrics['class_avg_iou'] = mean_shape_ious test_metrics['inctance_avg_iou'] = np.mean(all_shape_ious) print(xxxxxx)#看一共有多少个点的，没啥用 log_string('Accuracy is: %.5f'%test_metrics['accuracy']) log_string('Class avg accuracy is: %.5f'%test_metrics['class_avg_accuracy']) log_string('Class avg mIOU is: %.5f'%test_metrics['class_avg_iou']) log_string('Inctance avg mIOU is: %.5f'%test_metrics['inctance_avg_iou'])if __name__ == '__main__': args = parse_args() main(args)接着是有关文件处理（数据预处理等）的代码：
import osfilePath = "D:\\Learning\\无人机项目\\Pointnet2\\Pointnet2\\data\\shapenetcore_partanno_segmentation_benchmark_v0_normal\\03797390\\"#####最后一行会出现一个， 报错！！！！！！！！！######手动删除或改进程序#####file = '1.txt'with open(file,'a') as f: f.write("[") for i,j,k in os.walk(filePath): for name in k:base_name=os.path.splitext(name)[0] #去掉后缀 .txtf.write(" \"")f.write(os.path.join("shape_data/03797390/",base_name))f.write("\"")f.write(",") f.write("]")f.close()# -*- coding:utf-8 -*-import osfilePath = 'D:\\Learning\\无人机项目\\Pointnet2\\Pointnet2\\data\\shapenetcore_partanno_segmentation_benchmark_v0_normal\\02691156\\'for i,j,k in os.walk(filePath): for name in k: list1 = [] for line in open(filePath+name): a = line.split() #print(a) b = a[0:6] #print(b) a1 =float(a[0]) a2 =float(a[1]) a3 =float(a[2]) #print(a1) if(a1==0 and a2==0 and a3==0): continue list1.append(b[0:6]) with open(filePath+name, 'w+') as file: for i in list1: file.write(str(i[0])) file.write(' '+str(i[1])) file.write(' ' + str(i[2])) file.write(' ' + str(i[3])) file.write(' ' + str(i[4])) if(i!=list[-1]): file.write('\n') file.close()# print(list)# import os# filePath = '161865156110305.txt'# for i,j,k in os.walk(filePath):# for name in k:# print(name)# f = open(filePath+name) # 打开txt文件# line = f.readline() # 以行的形式进行读取文件# list1 = []# while line:# a = line.split()# b = a[0:3]# 这是选取需要读取/修改的列 前两列# c = float(a[-1])# a1 =float(a[0])# a2 =float(a[1])# a3 =float(a[2])# if(a1==0 and a2==0 and a3==0)# print(c)# if(float(a[-1])==36.0):# c=2# if(float(a[-1])==37.0):# c=3# b.append(c)# list1.append(b) # 将其添加在列表之中# line = f.readline()# f.close()# print(list1)# with open(filePath+name, 'w+') as file:# for i in list1:# file.write(str(i[0]))# file.write(' '+str(i[1]))# file.write(' ' + str(i[2]))# file.write(' ' + str(i[3]))# if(i!=list[-1]):#file.write('\n')# file.close()# path_out = 'test.txt' # 新的txt文件# with open(path_out, 'w+') as f_out:# for i in list1:# fir = '9443_' + i[0]# 第一列加前缀'9443_'# sec = 9443 + int(i[1]) # 第二列数值都加9443# # print(fir)# # print(str(sec))# f_out.write(fir + ' ' + str(sec) + '\n') # 把前两列写入新的txt文件祝小伙伴们心想事成！！！