1 はじめに
舌のセグメンテーションは、舌の診断と検出の基礎であり、舌を正確にセグメンテーションすることによってのみ、その後のトレーニングと予測の精度を確保できます。この部分の実際のタスクは、ユーザーがアップロードした画像で舌に属する正確なピクセルを見つけることです。舌セグメンテーションは、生物医学画像セグメンテーションの分野に属します。分割効果は次のとおりです。
2 データセットの紹介
舌画像データセットには、元の舌画像とセグメント化された 2 値画像が含まれており、合計 979 枚*2 枚の画像が含まれています。サンプル画像は次のとおりです。
データセット+ソースコード取得方法:
Xianyuリンク
【闲鱼】https://m.tb.cn/h.Ud0AYdZ?tk=40LV2COORtc CZ0001 「我在闲鱼发布了【舌象图片数据集 舌头图片数据集 舌头图片分割标注完成的数据集】」
点击链接直接打开
3 モデル紹介
U-Net は優れたセマンティック セグメンテーション モデルであり、Zhong-e 診断における U-Net の 3 つの部分、すなわち、主要な特徴抽出部分、強化された特徴抽出部分、および予測部分があります。主特徴抽出部分を使用して 5 つの予備的な有効特徴レイヤーを取得し、拡張特徴抽出部分を使用して、上記で取得した 5 つの有効特徴レイヤーをアップサンプリングし、特徴融合を実行します。最後に、すべての特徴を組み合わせた有効特徴レイヤーが取得され、最終的な有効特徴レイヤーを使用してピクセルを予測し、舌に属するピクセルを見つけます。特定の操作の詳細を次の図に示します。
ラベル付けの後、PyTorch フレームワークを使用して U-Net モデルが構築され、舌画像の特徴がキャプチャされ、舌画像のラベルが予測されました。モデルを評価するために、各反復の平均損失率がトレーニング中に計算されます。画像ごとの最終的な損失は約 2% です。具体的な平均損失率の変化は次のとおりです。
トレーニングは 4 日間続き、合計 979 個のラベル付き画像が使用され、最終的な平均予測損失率は約 2% でした。このモデルは、舌のセグメンテーションの効果が非常に理想的であると予測しています. これは、損失率が 40% であり、損失率が 2% である場合のセグメンテーション結果の例です. 例を次の図に示します:
(1) セグメンテーション損失率40%の場合の結果 図
(2)損失率2%の場合のセグメンテーション結果
モデルの予測結果より、舌に属するピクセルをマッチングして抽出し、舌に属さない部分を抽出します。最終的な舌セグメンテーション効果は次のとおりです。
4 コード実装の詳細
4.1 関連文書の紹介
notedata フォルダーにセグメンテーションと注釈の画像、ordata フォルダーに元の画像、params フォルダーにトレーニング モデル ファイル、結果フォルダーにテスト サンプル画像、train_image フォルダーにトレーニング プロセス画像があります。
4.2 ユーティリティ.py
ツール クラス: データセット内の各画像のサイズが異なるため、フォローアップ作業をスムーズに進めるために、ツール クラスをここで定義して、画像を 256*256 にスケーリングする必要があります (次のように変更できます)。あなた自身のニーズ)。
from PIL import Image
def keep_image_size_open(path, size=(256, 256)):
img = Image.open(path)
temp = max(img.size)
mask = Image.new('RGB', (temp, temp), (0,0,0))
mask.paste(img, (0,0))
mask = mask.resize(size)
return mask
4.3 データ.py
ここでの主な目的は、データセット内のラベル画像を元の画像と一致させてマージすることです〜特定の手順については、コード コメントに詳細な説明があります。
import os
from torch.utils.data import Dataset
from utils import *
from torchvision import transforms
transform = transforms.Compose([
transforms.ToTensor()
])
class MyDataset(Dataset):
def __init__(self, path): #拿到标签文件夹中图片的名字
self.path = path
self.name = os.listdir(os.path.join(path, 'notedata'))
def __len__(self): #计算标签文件中文件名的数量
return len(self.name)
def __getitem__(self, index): #将标签文件夹中的文件名在原图文件夹中进行匹配(由于标签是png的格式而原图是jpg所以需要进行一个转化)
segment_name = self.name[index] #XX.png
segment_path = os.path.join(self.path, 'notedata', segment_name)
image_path = os.path.join(self.path, 'ordata', segment_name.replace('png', 'jpg')) #png与jpg进行转化
segment_image = keep_image_size_open(segment_path) #等比例缩放
image = keep_image_size_open(image_path) #等比例缩放
return transform(image), transform(segment_image)
if __name__ == "__main__":
data = MyDataset("E:/ITEM_TIME/project/UNET/")
print(data[0][0].shape)
print(data[0][1].shape)
データセットが正則化されていることがわかります。
4.4 net.py
Unetネットワークの書き込み!
from torch import nn
import torch
from torch.nn import functional as F
class Conv_Block(nn.Module): #卷积
def __init__(self, in_channel, out_channel):
super(Conv_Block, self).__init__()
self.layer = nn.Sequential(
nn.Conv2d(in_channel, out_channel, 3, 1, 1, padding_mode='reflect',
bias=False),
nn.BatchNorm2d(out_channel),
nn.Dropout2d(0.3),
nn.LeakyReLU(),
nn.Conv2d(out_channel, out_channel, 3, 1, 1, padding_mode='reflect',
bias=False),
nn.BatchNorm2d(out_channel),
nn.Dropout2d(0.3),
nn.LeakyReLU()
)
def forward(self, x):
return self.layer(x)
class DownSample(nn.Module): #下采样
def __init__(self, channel):
super(DownSample, self).__init__()
self.layer = nn.Sequential(
nn.Conv2d(channel, channel,3,2,1,padding_mode='reflect',
bias=False),
nn.BatchNorm2d(channel),
nn.LeakyReLU()
)
def forward(self,x):
return self.layer(x)
class UpSample(nn.Module): #上采样(最邻近插值法)
def __init__(self, channel):
super(UpSample, self).__init__()
self.layer = nn.Conv2d(channel, channel//2,1,1)
def forward(self,x, feature_map):
up = F.interpolate(x, scale_factor=2, mode='nearest')
out = self.layer(up)
return torch.cat((out,feature_map),dim=1)
class UNet(nn.Module):
def __init__(self):
super(UNet, self).__init__()
self.c1=Conv_Block(3,64)
self.d1=DownSample(64)
self.c2=Conv_Block(64, 128)
self.d2=DownSample(128)
self.c3=Conv_Block(128,256)
self.d3=DownSample(256)
self.c4=Conv_Block(256,512)
self.d4=DownSample(512)
self.c5=Conv_Block(512,1024)
self.u1=UpSample(1024)
self.c6=Conv_Block(1024,512)
self.u2=UpSample(512)
self.c7=Conv_Block(512,256)
self.u3=UpSample(256)
self.c8=Conv_Block(256,128)
self.u4=UpSample(128)
self.c9=Conv_Block(128,64)
self.out = nn.Conv2d(64,3,3,1,1)
self.Th = nn.Sigmoid()
def forward(self,x):
R1 = self.c1(x)
R2 = self.c2(self.d1(R1))
R3 = self.c3(self.d2(R2))
R4 = self.c4(self.d3(R3))
R5 = self.c5(self.d4(R4))
O1 = self.c6(self.u1(R5,R4))
O2 = self.c7(self.u2(O1,R3))
O3 = self.c8(self.u3(O2,R2))
O4 = self.c9(self.u4(O3,R1))
return self.Th(self.out(O4))
if __name__ == "__main__":
x = torch.randn(2, 3, 256, 256)
net = UNet()
print(net(x).shape)
結果の一致は問題がないことを示しています〜
4.5 train.py
トレーニングコード〜
from torch import nn
from torch import optim
import torch
from data import *
from net import *
from torchvision.utils import save_image
from torch.utils.data import DataLoader
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
weight_path = 'params/unet.pth'
data_path = 'E:/ITEM_TIME/project/UNET/'
save_path = 'train_image'
if __name__ == "__main__":
dic = []###
data_loader = DataLoader(MyDataset(data_path),batch_size=3,shuffle=True) #batch_size用3/4都可以看电脑性能
net = UNet().to(device)
if os.path.exists(weight_path):
net.load_state_dict(torch.load(weight_path))
print('success load weight')
else:
print('not success load weight')
opt = optim.Adam(net.parameters())
loss_fun = nn.BCELoss()
epoch = 1
while True:
avg = []###
for i, (image,segment_image) in enumerate(data_loader):
image,segment_image = image.to(device),segment_image.to(device)
out_image = net(image)
train_loss = loss_fun(out_image, segment_image)
opt.zero_grad()
train_loss.backward()
opt.step()
if i%5 == 0:
print('{}-{}-train_loss===>>{}'.format(epoch,i,train_loss.item()))
if i%50 == 0:
torch.save(net.state_dict(), weight_path)
#为方便看效果将原图、标签图、训练图进行拼接
_image = image[0]
_segment_image = segment_image[0]
_out_image = out_image[0]
img = torch.stack([_image,_segment_image,_out_image],dim=0)
save_image(img, f'{save_path}/{i}.jpg')
avg.append(float(train_loss.item()))###
loss_avg = sum(avg)/len(avg)
dic.append(loss_avg)
epoch += 1
print(dic)
コードが正常に実行されていることがわかります。上記の損失率は 4 日間のトレーニング後の効果です。最初は非常に悪いに違いありません。忍耐が必要です。
4.6 テスト.py
画像をインテリジェントにセグメント化するコードをテストしてください~
from net import *
from utils import keep_image_size_open
import os
import torch
from data import *
from torchvision.utils import save_image
from PIL import Image
import numpy as np
net = UNet().cpu() #或者放在cuda上
weights = 'params/unet.pth' #导入网络
if os.path.exists(weights):
net.load_state_dict(torch.load(weights))
print('success')
else:
print('no loading')
_input = 'xxxx.jpg' #导入测试图片
img = keep_image_size_open(_input)
img_data = transform(img)
print(img_data.shape)
img_data = torch.unsqueeze(img_data, dim=0)
print(img_data)
out = net(img_data)
save_image(out, 'result/result.jpg')
save_image(img_data, 'result/orininal.jpg')
print(out)
#E:\ITEM_TIME\UNET\ordata\4292.jpg
img_after = Image.open(r"result\result.jpg")
img_before = Image.open(r"result\orininal.jpg")
#img.show()
img_after_array = np.array(img_after)#把图像转成数组格式img = np.asarray(image)
img_before_array = np.array(img_before)
shape_after = img_after_array.shape
shape_before = img_before_array.shape
print(shape_after,shape_before)
#将分隔好的图片进行对应像素点还原,即将黑白分隔图转化为有颜色的提取图
if shape_after == shape_before:
height = shape_after[0]
width = shape_after[1]
dst = np.zeros((height,width,3))
for h in range(0,height):
for w in range (0,width):
(b1,g1,r1) = img_after_array[h,w]
(b2,g2,r2) = img_before_array[h,w]
if (b1, g1, r1) <= (90, 90, 90):
img_before_array[h, w] = (144,238,144)
dst[h,w] = img_before_array[h,w]
img2 = Image.fromarray(np.uint8(dst))
img2.save(r"result\blend.png","png")
else:
print("失败!")
結果表示:
(1)元画像(orininal.jpg):(
2)モデルセグメンテーション画像(result.jpg):(
3)対応画素復元画像(blend.png):(2)の画像は白い部分が元の画像のピクセルで塗りつぶされ、黒い部分は緑で塗りつぶされます
この時点で、舌体のセグメンテーションは完了です。