使用Unet训练自己的模型

从github上下载代码：https://github.com/zhixuhao/unet制作自己的数据集1、将数据由24位深度改为8位：import osimport globimport cv2def to_gray(img,outdir):src = cv2.imread(img)try:dst = cv2.cvtColor(src,cv2.COLOR_BGR2GRAY)cv2.imwrit

文章共619字 · 阅读需要大约3分钟

一键AI生成摘要，助你高效阅读

问答

qq_41007606

6471人浏览 · 2020-05-16 13:03:15

qq_41007606 · 2020-05-16 13:03:15 发布

从github上下载代码：https://github.com/zhixuhao/unet

制作自己的数据集

1、将数据由24位深度改为8位：

import os
import glob
import cv2


def to_gray(img,outdir):
    src = cv2.imread(img)
    try:
        dst = cv2.cvtColor(src,cv2.COLOR_BGR2GRAY)
        cv2.imwrite(os.path.join(outdir,os.path.basename(img)), dst)
    except Exception as e:
        print(e)

for file in glob.glob('unet-master\\data\\membrane\\train\\label\\*.png'):
    to_gray(file,"unet-master\\data\\membrane\\train\\label1\\")

2、使用labelme.exe标注数据集

保存后会生成对应图片名的json文件：

将json文件批量转化为训练集：

import os
from glob import glob

json_files = glob(os.path.join(r"image", "*.json"))
for json_file in json_files:
    os.system(f"labelme_json_to_dataset {json_file} -o {''.join(json_file.split('.')[:-1])}")

将标注区域转化为白色，我的输入大小是96*512：

import os
import shutil
import cv2


path = r"image"
img_output = "unet-master\\data\\membrane\\train\\image"
label_output = "unet-master\\data\\membrane\\train\\label"

dirs = os.listdir(path)
for dir_name in dirs:
    filenames = os.listdir(os.path.join(path, dir_name))
    for filename in filenames:
        if filename == "img.png":
            img = cv2.imread(os.path.join(path, dir_name, filename), 0)
            img = cv2.resize(img, (96, 512))
            cv2.imwrite(os.path.join(img_output, dir_name+".png"), img)

        if filename == "label.png":
            img = cv2.imread(os.path.join(path, dir_name, filename), 0)
            img = cv2.resize(img, (96, 512))
            binary = cv2.threshold(img, 0, 255, cv2.THRESH_OTSU)[1]
            cv2.imwrite(os.path.join(label_output, dir_name+".png"), binary)

        print("Generate "+os.path.join(img_output, dir_name+".png"))

修改main.py中的输入尺寸：

from model import *
from data import *
import keras

#os.environ["CUDA_VISIBLE_DEVICES"] = "0"

data_gen_args = dict(rotation_range=0.2,
                    width_shift_range=0.05,
                    height_shift_range=0.05,
                    shear_range=0.05,
                    zoom_range=0.05,
                    horizontal_flip=True,
                    fill_mode='nearest')
myGene = trainGenerator(1,'data/membrane/train','image3','label3',data_gen_args,save_to_dir = None, target_size=(512, 96)) # 高， 宽

model = unet()
model_checkpoint = ModelCheckpoint('unet_membrane.hdf5', monitor='loss',verbose=1, save_best_only=True)
# reduce_lr = keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=10, verbose=0, mode='auto', epsilon=0.0001, cooldown=0, min_lr=0)

model.fit_generator(myGene,steps_per_epoch=60,epochs=10,callbacks=[model_checkpoint], shuffle=True)
model.save("./u-net-1.h5")

testGene = testGenerator("data/membrane/test", target_size=(512, 96)) # 高宽
results = model.predict_generator(testGene,30,verbose=1)
saveResult("data/membrane/test",results)

执行训练程序：

python main.py