点云实例分割-Softgroup-训练自己数据集

之前写的blog文章使用Softgroup官方自带的S3DIS配置文件，数据集格式和Dataset文件训练数据集，训练起来会方便些，但是自己训练数据集会处理这些操作，本文主要解决这个目的

• 这里我使用一个公共的植物点云数据集，训练Softgroup模型
• 公共数据集使用的是一个包含102个大豆的点云数据，其中对叶子和茎进行实例标注
• 主要步骤
  • 分析公共数据集，并通过脚本转化成S3DIS格式的数据集
  • 创建自己CustomDataset文件
  • 创建配置文件，训练模型
    

训练数据集和数据集处理

分析公共数据集

• 官方文档处理成的文件数据格式 https://github.com/thangvubk/SoftGroup/blob/main/docs/custom_dataset.md

• 下面展示的是下载下来的文件结构，下面根据文件的格式

• ![[Pasted image 20240813105842.png]]

• 每个标注的文件格式，共有6维分别前3个xyz三维坐标，后3个相应点的RGB信息
  

编写处理数据集程序

• 现在我们根据原始数据的结构和官方文档制作我们的训练模型的数据集
• 官方文档处理成的文件数据格式 https://github.com/thangvubk/SoftGroup/blob/main/docs/custom_dataset.md

import os
import numpy as np
import torch
import open3d as o3d
from pathlib import Path
from tqdm import tqdm
from glob import glob

# 加载和处理每个项目中的注释文件
def load_and_process_annotations(folder_path, classname_list):
    output = {
        "coord": [],  # 存储坐标
        "color": [],  # 存储颜色
        "semantic_gt": [],  # 存储语义标签
        "instance_gt": []   # 存储实例标签
    }

    annotations = glob(os.path.join(folder_path, '*'))
    for index, annotation in enumerate(annotations):
        stem = Path(annotation).stem
        classname, _ = stem.split('_')

        data = np.loadtxt(annotation)
        num_points = data.shape[0]

        sem = np.ones((num_points, 1), dtype=np.int8) * (classname_list.index(classname) + 1)
        ins = np.ones((num_points, 1), dtype=np.int8) * (index + 1)

        coord = data[:, :3]
        color = data[:, 3:6]

        output['coord'].append(coord)
        output['color'].append(color)
        output['semantic_gt'].append(sem)
        output['instance_gt'].append(ins)

    for key in output:
        output[key] = np.concatenate(output[key], axis=0)
    return output

# 对数据进行采样，以控制处理的点数
def resample_data(output, num_points_to_keep):
    total_points = output['coord'].shape[0]

    if total_points > num_points_to_keep:
        chosen_indices = np.random.choice(total_points, num_points_to_keep, replace=False)
    else:
        chosen_indices = np.arange(total_points)

    for key in output:
        output[key] = output[key][chosen_indices]

    return output

# 归一化颜色数据到 [-1, 1]
def normalize_color(output):
    data_min = output['color'].min(axis=0)
    data_max = output['color'].max(axis=0)
    output['color'] = (output['color'] - data_min) / (data_max - data_min)
    output['color'] = output['color'] * 2 - 1

    return output

# 主函数，遍历数据集中的所有项目，并进行处理
def main(folder_url, num_points_to_keep=10000):
    classname_list = ['leaf', 'stem', 'mainstem']
    projects = glob(os.path.join(folder_url, '*', 'Annotations'))
    
    for project in tqdm(projects):
        project_name = project.split('\\')[-2]
        output = load_and_process_annotations(project, classname_list)
        output = resample_data(output, num_points_to_keep)
        output = normalize_color(output)

        torch.save((output['coord'], output['color'], output['semantic_gt'].reshape(-1), output['instance_gt'].reshape(-1)), os.path.join('soybean-softgroup', f'{project_name}.pth'))


folder_url = r"D:\\dataset\\detection\\soybean3d\\"
main(folder_url)

生成数据集结果

• 这里我是手动划分的数据集，大家可以根据自己需求自己划分数据集，下面是最后的生成的目录结构

创建自己的CustomDataset文件和配置文件

CustomDataset文件

根据SoftGroup/softgroup/data/custom.py文件创建SoftGroup/softgroup/data/soybean.py文件，这里修改名称，类别和ID，根据我上面预处理数据集的顺序

把创建的文件类加入到SoftGroup/softgroup/data/__init__.py文件中，这里需要修改3处，根据原有的格式进行修改就行

配置文件

创建训练的配置文件/root/SoftGroup/configs/softgroup/softgroup_soybean.yaml，下面主要是根据自己数据集进行修改，还有训练参数大家根据自己电脑的性能进行修改就行

model:
  channels: 32
  num_blocks: 7
  semantic_classes: 4
  instance_classes: 3
  sem2ins_classes: []
  semantic_only: False
  ignore_label: -100
  grouping_cfg:
    score_thr: 0.2
    radius: 0.9
    mean_active: 300
    class_numpoint_mean: [-1, 1000, 1000,1000]
    npoint_thr: 0.05  # absolute if class_numpoint == -1, relative if class_numpoint != -1
    ignore_classes: [0]
  instance_voxel_cfg:
    scale: 3
    spatial_shape: 20
  train_cfg:
    max_proposal_num: 200
    pos_iou_thr: 0.5
  test_cfg:
    x4_split: False
    cls_score_thr: 0.001
    mask_score_thr: -0.5
    min_npoint: 100
    eval_tasks: ['semantic', 'instance']
  fixed_modules: []

data:
  train:
    type: 'soybean'
    data_root: '/root/soybean-softgroup'
    prefix: 'train'
    suffix: '.pth'
    repeat: 1
    training: True
    voxel_cfg:
      scale: 3
      spatial_shape: [128, 512]
      max_npoint: 250000
      min_npoint: 5000
  test:
    type: 'soybean'
    data_root: '/root/soybean-softgroup'
    prefix: 'val'
    suffix: '.pth'
    training: False
    voxel_cfg:
      scale: 3
      spatial_shape: [128, 512]
      max_npoint: 250000
      min_npoint: 5000

dataloader:
  train:
    batch_size: 8
    num_workers: 8
  test:
    batch_size: 1
    num_workers: 1

optimizer:
  type: 'Adam'
  lr: 0.001

fp16: False
epochs: 300
step_epoch: 0
save_freq: 5
pretrain: False
work_dir: ''

注意参数设置

• scale和radius参数设置
  • 根据 https://github.com/thangvubk/SoftGroup/blob/main/docs/custom_dataset.md 其中Config章节设置
  • 在Softgroup模型训练结果跟radius这个值影响很大，而且这个值不太固定，我训练自己数据的时候使用网格搜索找到的
  • 也可以
• data中参数设置
  • 例如删除 X4_split这是其他dataset单独处理的参数
• 请大家仔细检查参数，训练错误都会出现这里，还有就是预处理后的数据格式

训练模型

• 这里看到我们的数据集加载出来了，是我们数据集的个数
  

• 验证的结果
  

本文使用的程序和数据集数据

• 如果想完整的项目私信我
• 处理好的数据我上传到csdn上了 https://download.csdn.net/download/u012901740/89634574

后续

如果大家对点云分割感兴趣可以关注我，后续我会对三维重建和点云分割网络的改进进行更新