【语义分割】使用MMSegmentation训练VOC格式自定义数据集

原创已于 2023-12-06 06:45:02 修改 · 1.2k 阅读

12 ·

CC 4.0 BY-SA版权

文章标签：

#1024程序员节 #语义分割

于 2023-10-24 09:19:42 首次发布

Open MMLab 同时被 2 个专栏收录

5 篇文章

订阅专栏

语义分割

2 篇文章

订阅专栏

本文详细介绍了如何在MMSeg中准备VOC格式的数据集，选择并修改模型配置，如DeepLabV3，以及进行模型训练和图片批量推理的过程。同时提到了与MMDetection库相关的目标检测和自定义COCO数据集的训练示例。

部署运行你感兴趣的模型镜像

1.自定义数据集(VOC格式)的准备

1.1. 将数据集按照以下路径存放到MMSeg目录

mmsegmentation
|--data
|  |--VOCdevkit
|     |--VOC2012
|        |--Annotations
|           |-- 01.xml
|           |-- 02.xml
|
|        |--ImageSets
|           |-- Main
|           |-- Segmentation
|
|        |-- JPEGImages
|           |-- 01.jpg
|           |-- 02.jpg
|
|        |-- SegmentationClass
|           |-- 01.png
|           |-- 02.png|

1.2. 选择模型和预训练权重

1.3. 修改数据集配置文件

查看所选模型的配置文件，以训练DeepLabV3+为例：

(1) 修改数据集信息为：

'../_base_/datasets/pascal_voc12.py'

(2) 进入模型根目录文件修改classes数：configs/_base_/models/deeplabv3_r50-d8.py

(3) 进入voc数据集配置文件修改数据集类别信息：

(4) 进入训练参数配置文件修改预训练模型：configs/_base_/default_runtime.py

1.4. 每次修改完config文件需进行setup

python setup.py install

2. 模型训练

python tools/train.py config.py

以训练DeepLabV3+为例：

python tools/train.py configs/deeplabv3/deeplabv3_r50-d8_4xb4-40k_voc12aug-512x512.py

3. 模型图片批量推理

import os
import torch
import cv2
import argparse
import numpy as np
from pprint import pprint
from tqdm import tqdm
from mmseg.apis import init_model, inference_model

DEVICE = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
IMAGE_FILE_PATH = r"data/VOCdevkit/VOC2012/JPEGImages"
CONFIG = r'work_dirs/deeplabv3_r50-d8_4xb4-40k_voc12aug-512x512/deeplabv3_r50-d8_4xb4-40k_voc12aug-512x512.py'
CHECKPOINT = r'work_dirs/deeplabv3_r50-d8_4xb4-40k_voc12aug-512x512/iter_40000.pth'
SAVE_DIR = r"work_dir\infer_results"


def parse_args():
    parser = argparse.ArgumentParser(description='Visualize CAM')
    parser.add_argument('--img', default=IMAGE_FILE_PATH, help='Image file')
    parser.add_argument('--config', default=CONFIG, help='Config file')
    parser.add_argument('--checkpoint', default=CHECKPOINT, help='Checkpoint file')
    parser.add_argument('--device', default=DEVICE, help='device')
    parser.add_argument('--save_dir', default=SAVE_DIR, help='save_dir')

    args = parser.parse_args()
    return args


def make_full_path(root_list, root_path):
    file_full_path_list = []
    for filename in root_list:
        file_full_path = os.path.join(root_path, filename)
        file_full_path_list.append(file_full_path)
    return file_full_path_list


def read_filepath(root):
    from natsort import natsorted
    test_image_list = natsorted(os.listdir(root))
    test_image_full_path_list = make_full_path(test_image_list, root)
    return test_image_full_path_list


def apply_colormap(pred_mask):
    color_map = {
        0: [0, 0, 0],  # black: background
        1: [0, 255, 255],  # yellow: careful
        2: [255, 0, 0],  # blue: passable
        3: [0, 0, 255]  # red: impassable
    }

    rgb_image = np.zeros((*pred_mask.shape, 3), dtype=np.uint8)

    for key in color_map:
        rgb_image[pred_mask == key] = color_map[key]

    return rgb_image


def main():
    args = parse_args()

    model_mmseg = init_model(args.config, args.checkpoint, device=args.device)

    for imgs in tqdm(read_filepath(args.img)):
        result = inference_model(model_mmseg, imgs)
        pred_mask = result.pred_sem_seg.data.squeeze(0).detach().cpu().numpy().astype(np.uint8)

        colored_mask = apply_colormap(pred_mask)

        save_path = os.path.join(args.save_dir, f"{os.path.basename(args.config).split('.')[0]}")
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        cv2.imwrite(os.path.join(save_path, f"{os.path.basename(result.img_path).split('.')[0]}.png"), colored_mask,
                    [cv2.IMWRITE_PNG_COMPRESSION, 0])


if __name__ == '__main__':
    main()