DICOM画像を訓練用とテスト用に分けてみる
Keyword:python, DICOM， クロスバリデーション
AIの実験では，ある訓練データの集合に対して学習を行って，その学習モデルを使って未知のテストデータに対して推論を行い，その正答率から学習モデルの性能を評価します。
この実験のためにたくさん集めたDICOM画像を訓練用とテスト用に用手的に振り分けるのは大変な作業になります。
そこで，ここではラベル付されたDICOM画像を訓練用とテスト用に分けるプログラムを作成します。
つまり，画像ファイル名とラベル番号のリストになった訓練用CSVファイルとテスト用CSVファイルを作成するプログラムを作成します。
準備：ラベル付されたディレクトリの中にDICOM画像を格納する
このプログラムは，
「振り分けるDICOM画像は，親ディレクトリの下にラベル付されたサブディレクトリがあり，その中に該当する複数の画像が保存されている。」
こととします。
サンプルのDICOM画像はここにあります。私の全身のCT画像です。個人の実験にのみ使用し，他では使わないでください。
このサンプルDICOM画像では，部位を分類するタスクを考えます。
親ディレクトリdcmdir2の下に，Abdomen（腹部），Chest（胸部），Head（頭部）とLung（肺野）のラベルを持ったサブディレクトリがあるます。
その中にそれぞれの部位のDICOM画像が格納されています。（ラベル名に日本語使っても問題ありません）
ディレクトリの階層は次のようになります。
dcmdir2┬ Abdomen ┬ DICOM image a
│ │ ・・・
│ └ DICOM image b
├ Chest ┬ DICOM image c
│ │ ・・・
│ └ DICOM image d
├ Head ┬ DICOM image e
│ │ ・・・
│ └ DICOM image f
└ Lung ┬ DICOM image g
│ ・・・
└ DICOM image h
実際にはタスクに応じてラベルは変更します。
プログラムの考え方
学習を行うとき，ラベル名は0からはじまるラベル番号に置き換えます。
プログラムのはじめに，ラベル名とラベル番号を対応させたファイルlabelName.csvを作成します。画像をいくつのグループに分けるのか指定します。
例えば３つのグループ，0,1と2に分ける場合，この中から２つのグループを取り出して訓練データとしてTrainDataのDICOMファイル名のCSVファイルを作成します。残りのひとつをTestDataとしてDICOMファイル名のCSVファイルを作成します。
グループに分割する方法は最も単純な方法を考えました。
３つに分割するのであれば，3で割ったあまりが，0か１か２でそれぞれのグループに分けます。
結果として３グループに分ける場合，訓練用とテスト用のファイルが次のように作成されます。
TrainData.0.3.csv TestData0.3.csv
TrainData.１.3.csv TestData１.3.csv
TrainData.2.3.csv TestData2.3.csv
このCSVファイルには画像ファイル名，ラベル番号とラベル名がリストになっています。
プログラムを入力して実行してみよう
splitDicomData.py
import os
def saveLabelName(dataDir, saveDir):
labelStr = ""
labelName = {}
dlist = os.listdir(dataDir)
print("Subdirectory:",len(dlist),dlist)
for i, subdir in enumerate(dlist):
labelStr += str(i) + "," + subdir + "\n"
labelName[subdir] = str(i)
f = open(saveDir + "/labelName.csv", 'w')
f.write(labelStr)
f.close()
return labelName
def splitData(dataDir, nGrp, saveDir, target, labelname):
#リスト保存ファイル名
saveTrainFname = saveDir + "TrainData." + str(target) + "." + str(nGrp) + ".csv"
saveTestFname = saveDir + "TestData." + str(target) + "." + str(nGrp) + ".csv"
strTrain = "" #訓練画像のリストを保存する変数
strTest = "" #テスト画像のリストを保存する変数
dlist = os.listdir(dataDir) #カテゴリディレクトリの取得
for subdir in dlist:
flist = os.listdir(dataDir + "/" + subdir)
print("\nLabel:",subdir,len(flist),"images")
for j, fname in enumerate(flist):
if j % nGrp == target:
print(j,fname, " → Test")
strTest += dataDir + "/" + subdir + "/"+fname +","
strTest += labelName[subdir] + "," + subdir + "\n"
else:
print(j,fname, " → Train")
strTrain += dataDir + "/" + subdir + "/"+fname +","
strTrain += labelName[subdir] + "," + subdir + "\n"
#訓練データのリストファイルの保存
f = open(saveTrainFname,'w')
f.write(strTrain)
f.close()
f = open(saveTestFname,'w')
f.write(strTest)
f.close()
if __name__=='__main__':
dataDir = "../dcmdir2" #対象画像の親ディレクトリ名
nGroup = 3 #グループ数
saveDir = "./TrainTestList/" #Listファイル保存ディレクトリ名
labelName = saveLabelName(dataDir, saveDir)
for target in range(nGroup):
splitData(dataDir, nGroup, saveDir, target, labelName)
実行結果
プログラムを実行すると，はじめにラベル名とラベル番号のファイルが作成が行われます。ラベル付けされたサブディレクトリの４つが表示されます。
$ python splitDicomData.py Subdirectory: 4 ['Lung', 'Abdomen', 'Chest', 'Head']
結果として，labelName.csvがTrainTestListディレクトリ内に作成されます。
ファイルの内容は次のようにラベル番号とラベル名の対応表になっています。
続けてプログラムは，各サブディレクトリを読み，その中のDICOM画像を訓練用かテスト用か振り分けていきます。
Label: Lung 61 images 0 E5306265 → Test 1 W5305890 → Train 2 K5303593 → Train 3 M5305390 → Test 4 R5305640 → Train 5 T5304000 → Train 6 C5304421 → Test 7 A5304343 → Train 8 D5304468 → Train 9 O5303765 → Test ・・・・・ 58 J5305250 → Train 59 D5303265 → Train 60 T5305734 → Test Label: Abdomen 84 images 0 H5259375 → Test 1 F5258046 → Train 2 L5258328 → Train ・・・・・
結果として，訓練用DICOMファイルのCSVファイルTrainData.i.n.csvとテスト用のTestData.i.n.csvが作成されます。
ファイルは，DICOM画像ファイル名，ラベル番号とラベル名のリストが保存されています。
このようにDICOM画像をディレクトリから移動させることなく，対象とする訓練データ，テストデータをファイルにリストとして保存することで柔軟にAIの実験を行うことができます。
プログラムの概説
saveLabelName関数
この関数はラベル名とラベル番号の対応をCSVファイルlabelName.csvを作成します。２つの引数dataDirとsaveDirを持ちます。
- dataDirは，ラベル付けされたサブディレクトリの親ディレクトリです。
- saveDirはlabelName.csvをはじめ訓練・テスト用のDICOM画像リストのCSVファイルを格納するディレクトリ名です。
labelStrはCSVファイルにラベル番号とラベル名を書き出すための変数です。
labelNameはラベル名をキーにラベル番号を値にもつ辞書です。
これらを初期化しています。
labelStr = " "
labelName = { }
次にディレクトリdataDir内のディレクトリのリストをdlistに格納します。
dlist = os.listdir( dataDir )
次にfor文を使って，ディレクトリのリストdlistからひとつづつディレクトリ名をsubdirに取り出します。
それと0からの計数を変数iに入れ，ラベル番号（i）とラベル名（substr）のCSVデータを1行分作成しいます。
labelStr += str( i ) + " , " + subdir + " \n "
同時にラベル名（subdir）をキーにして，ラベル番号（i）を値に辞書lableNameに追加しています。
labelName[subdir] = str( i )
最後にCSVデータをファイルに書き出します。
f = open( saveDir + " /labelName.csv ", 'w ')
f.write( labelStr )
f.close( )
辞書データlabelNameを関数の返却値としてreturnしています。
この辞書は，次のsplitData関数で使います。
splitData関数
この関数は画像ファイルを訓練用，テスト用に分けます。５つの引数を持ちます。
- dataDirは，DICOM画像が保存されている親ディレクトリ名です。
- nGrp,は，分割するグループ数です。
- saveDirは，訓練用・テスト用のDICOM画像リストファイルの保存ディレクトリ名です。
- targetは，テスト用の分割グループ番号です。指定された番号のグループがテスト用に，それ以外が訓練用にわけられます。
- labelnameは，ラベル名をキーにラベル番号を値に持つ辞書です。saveLabelName関数で作成されたものです。
ターゲットtargetとグループ番号nGrpから訓練用ファイル名saveTrainFnameとテスト用のファイル名saveTestFnameを作成しています。
訓練用とテスト用のCSVデータを作成するstrTrainとstrTestを初期化しています。
次に親ディレクトリdataDir内のサブディレクトリのリストをdlistに格納します。
for文によりひとつひとつディレクトリ名を取り出し，次の処理を行います。
取り出したひとつのサブディレクトリsubdir内のDICOM画像ファイルのリストflistを作成します。
flist = os.listdir( dataDir + " / " + subdir )
このファイルリストflistからfor文を用いて1画像づつ取り出し，その順序jも取り出します。
順序jが分割グールで割ったあまりが，もしターゲットtargetに等しければ，その画像をテスト用のCSVデータの変数に追加します。
if j % nGrp == target :
strTest += dataDir + " / " + subdir + " / "+fname +" , "
さらに，ラベル番号を辞書labelNameを用いて取得して追加し，ラベル名subdirの値を追加します。
もし，順序jが分割グールで割ったあまりが，ターゲットtargetに等しけなければ，その画像を訓練用のCSVデータの変数に追加します。
最後に訓練用・テスト用のCSVデータをそれぞれのCSVファイルに保存して終了です。
main関数
この関数がプログラム実行時に最初に呼ばれ，- 対象画像が入った親ディレクトリ
- 分割グループ数
- 各CSVファイルを保存するディレクトリ
が設定されています。
次にsaveLabelName関数を読んで，ラベル名の辞書をlabelNameに代入します。
グループ数分ループしてsplitData関数を呼び出して，訓練用・テスト用の画像リストのCSVファイルを作成します。
おわりに
このプログラムにより，簡単に訓練用・テスト用のDICOM画像のリストをCSVファイルに出力することができました。また，このCSVファイルを編集することにより，容易に訓練用データやテスト用データを修正することができます。
次回，Pyゼミ3.04では今回作成したlabelName.csvや各訓練・テストCSVファイルを用いてCT画像の部位の分類を行ってみたいと思います。
課 題
今回，n枚の画像を分類するのに分割数mで割ったあまりを見て判断したが，乱数を用いてランダムにm個に分割するようにrandomSplitData関数を作成しなさい。ただし，各グループの訓練用とテスト用にデータの重複があってはならない。
0 件のコメント:
コメントを投稿