當前位置：首頁 > 编程语言 > python >内容正文

python

【Python】mmSegmentation语义分割框架教程（自定义数据集、训练设定、数据增强）

發(fā)布時間：2025/3/11 python 17 豆豆

生活随笔收集整理的這篇文章主要介紹了【Python】mmSegmentation语义分割框架教程（自定义数据集、训练设定、数据增强）小編覺得挺不錯的,現(xiàn)在分享給大家,幫大家做個參考.

文章目錄

0.mmSegmentation介紹
1.mmSegmentation基本框架
- 1.1.mmSegmentation的model設置
- 1.2.mmSegmentation的dataset設置
- - 1.2.1.Dataset Class文件配置
  - 1.2.2.Dataset Config文件配置
  - 1.2.3.Total Config文件配置
2.運行代碼
3.展示效果圖和預測
X.附錄
- X.1.mmSegmentation框架解釋
- X.2.mmsegmentation使用的預訓練backbone
- X.2.mmsegmentation官方幫助文檔

0.mmSegmentation介紹

$\qquad$ mmSegmentation是openmmlab項目下開源的圖像語義分割框架，目前支持pytorch，由于其擁有pipeline加速，完善的數(shù)據(jù)增強體系，完善的模型庫，作為大數(shù)據(jù)語義分割訓練及測試的代碼框架是再好不過了。
$\qquad$ 在開始本教程之前，你需要解決openmmlab的環(huán)境配置問題，好在這個repo上已經(jīng)有很人性化的步驟講解了，在此附上鏈接，就不贅述了：

Github鏈接：安裝openmmlab環(huán)境

使用教程的相關鏈接如下（github的項目還自帶了中文版）：

Github鏈接：openmmlab/mmSegmentation
Gitio教程：openmmlab/mmSegmenatation

$\qquad$ 對著mmSegmentation官方教程一步步做固然是能做出來，但是由于其框架結構過于復雜，加之官方教程對如何規(guī)范自定義數(shù)據(jù)集缺乏一些tips，因而本文提供了一個相對簡單的教程供大家參考。本文所有講解目錄均為mmSegmentation的項目目錄。

1.mmSegmentation基本框架

$\qquad$ 要說mmSegmentation(以下簡稱mmSeg)當中最重要的東西，固然是Config文件了，Config文件可以分為4大類：

model config

dataset config

runtime config

schedule config

$\qquad$ 如果你想知道為什么分成這四大類，請參考本文X.1.節(jié)，對這個不感興趣就繼續(xù)往下看。其實3和4大多數(shù)人都用不到的，重點還是在1和2，下面就從這兩個角度給大家來一個不算精細的講解。

1.1.mmSegmentation的model設置

$\qquad$ 如果采用的是mmSegmentation里面支持的模型，那么固然是不需要自己寫class了，自己挑一個模型就可以了。這些model的目錄保存在了configs/models里面了。

第一個下劃線前面的都好理解，就是模型的名字唄，那r50-d8可能就是resnet的類型了，有人會問，那resnet101和resnet152哪去了，別急，其實這些只是baseline，它的backbone是可以改的，比如說我們要使用的是danet_r50-d8.py，我們先打開它（這里我已經(jīng)將SyncBN改成了BN，因為需要單GPU訓練）：

$\qquad$ 只需要把model.backbone.depth設為101或者152就可以使用resnet101或者resnet152啦，如果你的本地沒有模型，mmSeg就會從model_zoo里面下載一個，如果本地有（應該是保存在了checkpoint里面），則自動加載本地的，不會重復下載。其他的操作后面會講，另外如果你是多GPU操作就選擇使用SyncBN，否則就使用BN就可以了。如果使用了SyncBN卻只有一塊可用的GPU，那可能會報類似AssertionError:Default process group is not initialized的錯誤。有人可能問那我直接改了這個文件不就吧原來的默認參數(shù)給覆蓋了嘛，不要緊，看到后面大家就會明白這個問題很容易解決，這里只是給大家做一個demo。

1.2.mmSegmentation的dataset設置

$\qquad$ 數(shù)據(jù)集設置比model的稍微復雜一點，這里會直接定義一個自己的數(shù)據(jù)集（Custom Dataset）來說明其原理。數(shù)據(jù)集需要準備的文件有三個

Dataset Class文件

Dataset Config文件

Total Config文件

$\qquad$ 在X.1.節(jié)提到的config文件就是Total config（頂層設置文件），也是train.py文件直接調(diào)用的config文件，而Dataset Class文件是用來定義數(shù)據(jù)集的類別數(shù)和標簽名稱的，Dataset Config文件則是用來定義數(shù)據(jù)集目錄、數(shù)據(jù)集信息（例如圖片大小）、數(shù)據(jù)增強操作以及pipeline的。

1.2.1.Dataset Class文件配置

$\qquad$ 首先來說Dataset Class文件，這個文件存放在 mmseg/datasets/ 目錄下，

$\qquad$ 在這個目錄下自己建一個數(shù)據(jù)集文件，并命個名。配置文件實際上是繼承該目錄下custom.py當中的CustomDataset父類的，這樣寫起了就簡單多了，大多數(shù)情況下（當你的數(shù)據(jù)集是以一張張圖片出現(xiàn)并且可用PIL模塊讀入時），你只需要設置兩個參數(shù)即可——類別標簽名稱（CLASSES）和類別標簽上色的RGB顏色（PALETTE）。以我的配置文件為例，代碼如下：

from mmseg.datasets.builder import DATASETS from mmseg.datasets.custom import CustomDataset import os.path as osp@DATASETS.register_module() class MRDDataset(CustomDataset):CLASSES = ("background","road")PALETTE = [[0,0,0],[255,255,255]]def __init__(self, split, **kwargs):super().__init__(img_suffix='.png', seg_map_suffix='.png', split=split, **kwargs)assert osp.exists(self.img_dir) and self.split is not None

$\qquad$ img_suffix和seg_map_suffix分別是你的數(shù)據(jù)集圖片的后綴和標簽圖片的后綴，因個人差異而定，tif格式的圖片我還沒有試過，但是jpg和png的肯定是可以的。
$\qquad$ 設置好之后記得保存在mmseg/datasets/目錄下（我的文件名叫my_road_detect.py）。另外還需要設置一下該目錄下的__init__文件：

from .ade import ADE20KDataset from .builder import DATASETS, PIPELINES, build_dataloader, build_dataset from .chase_db1 import ChaseDB1Dataset from .cityscapes import CityscapesDataset from .custom import CustomDataset from .dataset_wrappers import ConcatDataset, RepeatDataset from .drive import DRIVEDataset from .hrf import HRFDataset from .pascal_context import PascalContextDataset, PascalContextDataset59 from .stare import STAREDataset from .voc import PascalVOCDataset from .my_road_detect import MRDDataset __all__ = ['CustomDataset', 'build_dataloader', 'ConcatDataset', 'RepeatDataset','DATASETS', 'build_dataset', 'PIPELINES', 'CityscapesDataset','PascalVOCDataset', 'ADE20KDataset', 'PascalContextDataset','PascalContextDataset59', 'ChaseDB1Dataset', 'DRIVEDataset', 'HRFDataset','STAREDataset',"MRDDataset" ]

$\qquad$ 需要改兩個地方，①import的時候要把自己的Dataset加載進來，②__all__數(shù)組里面需要加入自己的Dataset類名稱，修改完成之后保存。這兩部操作完成之后還不行，由于訓練的時候需要txt文件指示訓練集、驗證集和測試集的txt文件，一開始我以為這只是一個optional option，但無奈Custom Dataset的__init___下面給我來了一句assert osp.exists(self.img_dir) and self.split is not None，那好吧，不知道刪了and后面的條件會有什么后果，還是自己創(chuàng)一個吧，寫來一個簡單的劃分數(shù)據(jù)集并保存到txt的demo，大家可以把這個py文件放到你的數(shù)據(jù)集上一級目錄上并對著稍微改改：

import mmcv import os.path as osp data_root = "/data3/datasets/Custom/Lab/Segmentation/" ann_dir = "ann_png1" split_dir = 'splits' mmcv.mkdir_or_exist(osp.join(data_root, split_dir)) filename_list = [osp.splitext(filename)[0] for filename in mmcv.scandir(osp.join(data_root, ann_dir), suffix='.png')] with open(osp.join(data_root, split_dir, 'train.txt'), 'w') as f:# select first 4/5 as train settrain_length = int(len(filename_list)*4/5)f.writelines(line + '\n' for line in filename_list[:train_length]) with open(osp.join(data_root, split_dir, 'val.txt'), 'w') as f:# select last 1/5 as train setf.writelines(line + '\n' for line in filename_list[train_length:])

data_root寫自己的工作目錄名稱，ann_dir寫標簽圖片所在的目錄，split_dir則是在data_root下生成split txt文件保存的文件夾目錄，其他的就不需要怎么改了。如果你在data_root/split_dir/下成功找到了train.txt和val.txt文件，就沒有問題了。

1.2.2.Dataset Config文件配置

$\qquad$ Dataset Config文件在 configs/__base__/ 目錄下，需要自己新建一個xxx.py文件。

還是以我自己的Custom Dataset為例，它的書寫格式如下：

# dataset settings dataset_type = 'MRDDataset' data_root = '/data3/datasets/Custom/Lab/Segmentation/' img_norm_cfg = dict(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True) crop_size = (640, 480) train_pipeline = [dict(type='LoadImageFromFile'),dict(type='LoadAnnotations'),dict(type='Resize', img_scale=(640, 480), keep_ratio=True),dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),dict(type='RandomFlip', prob=0.5),dict(type='PhotoMetricDistortion'),dict(type='Normalize', **img_norm_cfg),dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),dict(type='DefaultFormatBundle'),dict(type='Collect', keys=['img', 'gt_semantic_seg']), ] test_pipeline = [dict(type='LoadImageFromFile'),dict(type='MultiScaleFlipAug',img_scale=(640, 480),# img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],flip=False,transforms=[dict(type='Resize', keep_ratio=True),dict(type='RandomFlip'),dict(type='Normalize', **img_norm_cfg),dict(type='ImageToTensor', keys=['img']),dict(type='Collect', keys=['img']),]) ] data = dict(samples_per_gpu=2,workers_per_gpu=2,train=dict(type=dataset_type,data_root=data_root,img_dir='data1_for_ann',ann_dir='ann_png1/',pipeline=train_pipeline,split="splits/train.txt"),val=dict(type=dataset_type,data_root=data_root,img_dir='data1_for_ann',ann_dir='ann_png1',split="splits/val.txt",pipeline=test_pipeline),test=dict(type=dataset_type,data_root=data_root,img_dir='data1_for_ann',ann_dir='ann_png1',split="splits/val.txt",pipeline=test_pipeline))

需要改的地方有以下幾個：

img_norm_cfg：數(shù)據(jù)集的方差和均值

crop_size：數(shù)據(jù)增強時裁剪的大小. img_dir：

img_scale：原圖像尺寸

data_root：工作目錄

split：之前操作做txt文件的目錄

sample_per_gpu：batch size

workers_per_gpu：dataloader的線程數(shù)目，一般設2，4，8，根據(jù)CPU核數(shù)確定,或使用os.cpu_count()函數(shù)代替

PhotoMetricDistortion是數(shù)據(jù)增強操作，有四個參數(shù)（參考博客）分別是亮度、對比度、飽和度和色調(diào)，它們的默認設定如下：

brightness_delta=32; # 32 contrast_range=(0.5, 1.5); # (0.5, 1.5)，下限-上限 saturation_range=(0.5, 1.5); # (0.5, 1.5)，下限-上限 hue_delta=18; # 18

如果不想使用默認設定，仿照其他選項將自定義參數(shù)寫在后面即可，例如

dict(type='PhotoMetricDistortion'，contrast_range=(0.5, 1.0))

改好之后保存 configs/__base__/ 目錄下。
$\qquad$ 這里也給大家提供了計算數(shù)據(jù)集方差和均值的一個樣例程序（多數(shù)據(jù)集計算整體均值和標準差）：

# -*- coding: utf-8 -*- """ Created on Fri Jun 25 10:38:17 2021@author: 17478 """ import os import cv2 import numpy as np from tqdm import tqdm # pip install tqdm import argparsedef input_args():parser = argparse.ArgumentParser(description="calculating mean and std")parser.add_argument("--data_fmt",type=str,default='samples_{name}')parser.add_argument("--data-name",type=str,nargs="+",default=['morning','noon','afternoon','dusk','snowy'])return parser.parse_args()if __name__ == "__main__":opt = input_args()img_files =[]for name in opt.data_name:img_dir = opt.data_fmt.format(name=name)files = os.listdir(img_dir)img_files.extend([os.path.join(img_dir,file) for file in files])meanRGB = np.asarray([0,0,0],dtype=np.float64)varRGB = np.asarray([0,0,0],dtype=np.float64)for img_file in tqdm(img_files,desc="calculating mean",mininterval=0.1):img = cv2.imread(img_file,-1)meanRGB[0] += np.mean(img[:,:,0])/255.0meanRGB[1] += np.mean(img[:,:,1])/255.0meanRGB[2] += np.mean(img[:,:,2])/255.0meanRGB = meanRGB/len(img_files)for img_file in tqdm(img_files,desc="calculating var",mininterval=0.1):img = cv2.imread(img_file,-1)varRGB[0] += np.sqrt(np.mean((img[:,:,0]/255.0-meanRGB[0])**2))varRGB[1] += np.sqrt(np.mean((img[:,:,1]/255.0-meanRGB[1])**2))varRGB[2] += np.sqrt(np.mean((img[:,:,2]/255.0-meanRGB[2])**2))varRGB = varRGB/len(img_files)print("meanRGB:{}".format(meanRGB))print("stdRGB:{}".format(varRGB))

1.2.3.Total Config文件配置

$\qquad$ Total Config文件是train.py直接調(diào)用的config文件，在X.1.節(jié)也有介紹，在此只說明如何即可。該文件在 config/xxxmodel/ 的目錄下，你選用的是哪一個model，就選擇哪一個目錄。

以DANet為例，我們書寫一個total config文件，并保存在configs/danet的文件夾下：

_base_ = ['../_base_/models/danet_r50-d8.py', '../_base_/datasets/my_road_detect.py','../_base_/default_runtime.py', '../_base_/schedules/schedule_20k.py' ] model = dict(decode_head=dict(num_classes=2),auxiliary_head=dict(num_classes=2))

$\qquad$ 這個代碼就一個__base__的數(shù)組，第一個元素代表模型路徑，也就是在1.1.節(jié)介紹的模型文件（在這個教程里就不帶著大家重寫模型了）；第二個元素代表數(shù)據(jù)集的Dataset config文件（詳見1.2.2節(jié)）；第三個元素和第四個元素本教程未涉及到，按照默認參數(shù)寫也沒有太大問題，如果想修改訓練的代數(shù)以及l(fā)og和save的頻率修改第4元素及響應文件，在此就不再贅述了。另外如果你的模型不是19類的（因為是原模型是根據(jù)cityscapes寫的，輸出通道為19），需按照上面修改一下。
$\qquad$ 到此為止要恭喜大家，代碼終于可以試跑了，如果你的代碼出現(xiàn)Error或者Exception也不要慌，從環(huán)境配置到流程一一對照一遍，調(diào)試大項目要有耐心，也歡迎大家評論區(qū)留言。

2.運行代碼

$\qquad$ 在項目目錄下，輸入python tools/train.py xxxconfig.py --work-dir=xxx即可運行，其中xxxconfig.py就是我們剛剛保存的Total config文件（記得要把完整路徑也加上），work-dir其實就是保存log和model的目錄（如果沒有會自己創(chuàng)建）。如果發(fā)現(xiàn)import mmseg找不到這個包，那八成是調(diào)試器運行目錄不在根目錄下造成的，要不就配置run的目錄，要不就直接吧tools/train.py復制到根目錄下運行。運行結果差不多是這樣：

使用gpustat的包查看gpu狀態(tài)

$\qquad$ 雖然我的數(shù)據(jù)集很小（做測試的，就50張圖片），但是gpu利用率仍然接近100%，可見其代碼優(yōu)化做的已經(jīng)相當理想了。（我開了NVIDIA的圖形加速，所以出現(xiàn)了很多其他的利用進程）。
$\qquad$ 這里有讀者會疑問為什么上面不顯示epoch，因為mmseg默認是iteration-based的，所謂iteration即batch的個數(shù)，若要改成epoch，則需要參考docs/config.md進行修改：

runner = dict(type='EpochBasedRunner',max_epoch='200') checkpoint_config = dict(by_epoch=True,interval=20) # save checkpoint per 20 epochs

以上代碼可放在Total config文件中。

3.展示效果圖和預測

$\qquad$ 最后寫了展示預測效果的代碼，把config_file和checkpoint_file替換成你自己的config文件和pth文件（保存模型的）即可：

from mmseg.apis import inference_segmentor, init_segmentor, show_result_pyplot from mmseg.core.evaluation import get_palette config_file = "configs/danet/danet_r50-d8_360x480_20k_mrd.py" checkpoint_file = 'work_dirs/danet_r50-d8_375x1242_20k_mrd/latest.pth' model = init_segmentor(config_file, checkpoint_file, device='cuda:0') img = '/data3/datasets/Custom/Lab/Segmentation/data1_for_ann/000000.png' result = inference_segmentor(model, img) show_result_pyplot(model, img, result, [[0,255,0],[255,255,255]])

$\qquad$ 我上的是白色（道路）和綠色（非道路），不是特別好看，哈哈，但是mask和img的相對位置很容易看出來，這個配顏色的話，大家還是自己定吧。我這個數(shù)據(jù)集太少，只是給大家做個演示，結果肯定是過擬合的。

X.附錄

X.1.mmSegmentation框架解釋

在mmSegmentation的項目目錄下，打開Configs/下面的目錄

隨便打開一個文件就知道了

從文件的名字也可以看出，它是模型（baseline+backbone、數(shù)據(jù)集、schedule的組合（runtime是default設置，就沒包含在名稱內(nèi)）。

X.2.mmsegmentation使用的預訓練backbone

預訓練backbone下載鏈接為：
mmcv預訓練模型下載地址（.json文件，復制對應模型的鏈接即可下載）

X.2.mmsegmentation官方幫助文檔

可在docs/tutorials中查看

希望本文對您有幫助，謝謝閱讀！

總結

以上是生活随笔為你收集整理的【Python】mmSegmentation语义分割框架教程（自定义数据集、训练设定、数据增强）的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。

上一篇： Linux(CentOS)安装apach
下一篇： HTML5-特效