BiSeNet V2源码实现分割任务

小杨~~~~

已于 2023-08-08 15:24:49 修改

阅读量630

点赞数 3

文章标签：深度学习人工智能机器学习

于 2023-08-08 15:19:37 首次发布

本文链接：https://blog.csdn.net/Yangxiaoaijiejie/article/details/132167145

版权

1.3.2 生成train.txt和val.txt文件

1.1 环境准备

我的电脑GPU信息查看：根据提供的信息，可以看出系统中有一个 GPU设备。在输出的表格中，只有一行显示了 GPU 的相关信息，其中 GPU 的索引为 0，名称为 "NVIDIA Graphics..."，后面训练的时候会用到设置环境变量，所以先提前查一下gpu信息。

nvidia-smi

1.2 模型推理

配置好环境后，先来验证一下环境有没有问题。

# 1、测试单张图片
python tools/demo.py --config configs/bisenetv2_city.py --weight-path /path/to/your/weights.pth --img-path ./example.png
# 比如我的命令：
python tools/demo.py --config configs/bisenetv2_city.py --weight-path weights/model_final_v2_city.pth --img-path example.png

# 2、测试视频
python tools/demo_video.py --config configs/bisenetv2_coco.py --weight-path res/model_final.pth --input ./video.mp4 --output res.mp4
# 比如我的命令：
python tools/demo_video.py --config configs/bisenetv2_city.py --weight-path weights/model_final_v2_city.pth --input test/video.mp4 --output res.mp4
# 注意：视频帧宽高太大可能会导致会报错

一些测试效果图：

1.3 准备cocostuff数据集

1.3.1 数据集下载并拆分

cocostuff语义分割数据集下载链接：COCO - Common Objects in Context (cocodataset.org)

1.3.2 生成train.txt和val.txt文件

（注意：一定记得检查labels、images文件夹下train2017和val2017是否有隐藏的文件（比如我就有.ipynb_checkpoints文件），如果有先删除再执行下列命令，否则会把隐藏文件一块写入新生成的txt文件内）

cd file
ls -a # 查看是否有隐藏的文件
rm -r .ipynb_checkpoints  # 删除该.ipynb_checkpoints隐藏文件
python tools/gen_dataset_annos.py --dataset coco  # 生成train.txt和val.txt两注释文件

新生成的俩txt文件：

1.3.3 检查数据集

# 统一命令：
python tools/check_dataset_info.py --im_root /path/to/your/data_root --im_anns /path/to/your/anno_file
# 比如我的命令：
python tools/check_dataset_info.py --im_root datasets/coco --im_anns datasets/coco/train.txt

# 如果查出来数据集不是171类，而是201类，则可以修改bisenetv2_coco.py文件的下列参数

n_cats=202, # 171

1.4 训练cocostuff数据集

提示：可以修改bisenetv2_coco.py文件的最大迭代次数参数：max_iter=600, #180000

# 训练命令方法一：
# 设置 CUDA_VISIBLE_DEVICES 环境变量；对于单块显卡，将其设置为索引为 0 的显卡
export CUDA_VISIBLE_DEVICES=0  
# 设置可用的显卡数量；在单块显卡的情况下，将其设置为 1
NGPUS=1
cfg_file=configs/bisenetv2_coco.py
torchrun --nproc_per_node=$NGPUS tools/train_amp.py --config $cfg_file

# 训练命令方法二：
torchrun --nproc_per_node=1 tools/train_amp.py --config configs/bisenetv2_coco.py

1.5 定义自己的数据集

如果要想用自己的数据集去做训练，应该把数据集按照下列格式做划分，并在gen_dataset_annos.py脚本中增加下列自定义数据集函数模块，再运行该脚本生成train.txt和val.txt文件。

def gen_custom():  # 自定义数据集
    '''
        root_path:
            |- images
                |- train
                |- val
            |- labels
                |- train
                |- val
    '''
    root_path = './datasets/custom'
    save_path = './datasets/custom/'
    for mode in ('train', 'val'):
        im_root = osp.join(root_path, f'images/{mode}')  # root_path/images/{mode}
        lb_root = osp.join(root_path, f'labels/{mode}')

        ims = os.listdir(im_root)
        lbs = os.listdir(lb_root)

        print(len(ims))
        print(len(lbs))

        im_names = [el.replace('.jpg', '') for el in ims]  # 去掉后缀名
        lb_names = [el.replace('.png', '') for el in lbs]
        common_names = list(set(im_names) & set(lb_names))

        lines = [
            f'images/{mode}/{name}.jpg,labels/{mode}/{name}.png'
            for name in common_names
        ]

        with open(f'{save_path}/{mode}.txt', 'w') as fw:
            fw.write('\n'.join(lines))
if __name__ == '__main__':
    parse = argparse.ArgumentParser()
    parse.add_argument('--dataset', dest='dataset', type=str, default='coco')
    args = parse.parse_args()

    if args.dataset == 'coco':
        gen_coco()
    elif args.dataset == 'ade20k':
        gen_ade20k()
    elif args.dataset == 'custom':
        gen_custom()

cd file
ls -a # 查看是否有隐藏的文件
rm -r .ipynb_checkpoints  # 删除该.ipynb_checkpoints隐藏文件
python tools/gen_dataset_annos.py --dataset custom  # 生成train.txt和val.txt两注释文件

# 检查数据集
python tools/check_dataset_info.py --im_root datasets/custom --im_anns datasets/custom/train.txt

# 根据检查结果，修改bisenetv2_custom.py文件的下列参数
n_cats=133, # 171

1.6 训练自己的数据集

# 训练命令方法一：
export CUDA_VISIBLE_DEVICES=0  
NGPUS=1
cfg_file=configs/bisenetv2_coco.py
torchrun --nproc_per_node=$NGPUS tools/train_amp.py --config $cfg_file
# 训练命令方法二：
torchrun --nproc_per_node=1 tools/train_amp.py --config configs/bisenetv2_coco.py