Tensorflow2.0 使用Kera 迭代器加载图像分割训练集

吊炸天的小鸿鸿

已于 2022-05-12 12:40:05 修改

阅读量572

点赞数

分类专栏：医学图像分割 AI python 文章标签： keras 人工智能深度学习

于 2022-05-12 12:16:39 首次发布

本文链接：https://blog.csdn.net/qq_36340554/article/details/124728478

版权

AI 同时被 3 个专栏收录

10 篇文章 1 订阅

订阅专栏

python

8 篇文章 2 订阅

订阅专栏

医学图像分割

4 篇文章 0 订阅

订阅专栏

本文介绍了如何在处理大数据时，通过继承自Keras的Sequence类，定义DataGenerator来实现分批次加载数据。重点展示了如何从图片和标签文件中读取并转换为适合模型训练的数据格式，适用于深度学习任务如U-Net模型训练。

摘要由CSDN通过智能技术生成

当遇到大数据时，无法将数据全部加载进内存，需要用到分批次加载，网上的方法很多都是关于分类数据集，记录一下分割数据集使用迭代器进行数据加载的方式。
主要从keras.utils.Sequence 继承后定义一个数据加载器 DataGenerator。
！！！注：本文的代码只展现了关键部分，不是完整代码

定义数据生成器

import glob
import tensorflow as tf
from model import unet
from tensorflow import keras
import math
import os
import cv2
import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense


class DataGenerator(keras.utils.Sequence):

    def __init__(self, data_img, data_mask, batch_size=1, shuffle=True):
        self.batch_size = batch_size
        self.data_img = data_img
        self.data_mask = data_mask
        self.indexes = np.arange(len(self.data_img))
        self.shuffle = shuffle

    def __len__(self):
        # 计算每一个epoch的迭代次数
        return math.ceil(len(self.data_img) / float(self.batch_size))

    def __getitem__(self, index):
        # 生成每个batch数据，这里就根据自己对数据的读取方式进行发挥了
        # 生成batch_size个索引
        batch_indexs = self.indexes[index * self.batch_size:(index + 1) * self.batch_size]
        # 根据索引获取datas集合中的数据
        batch_data_img = [self.data_img[k] for k in batch_indexs]
        batch_data_mask = [self.data_mask[k] for k in batch_indexs]

        # 生成数据
        X, y = self.data_generation(batch_data_img, batch_data_mask)

        return X, y

    def on_epoch_end(self):
        # 在每一次epoch结束是否需要进行一次随机，重新随机一下index
        if self.shuffle == True:
            np.random.shuffle(self.indexes)

    def data_generation(self, batch_data_img, batch_data_mask):
        images = []
        masks = []

        # 生成数据
        for data_img, data_mask in zip(batch_data_img, batch_data_mask):
            # x_train数据
            image = cv2.imread(data_img,cv2.IMREAD_COLOR)
            image = cv2.resize(image,(256,256))

            image = list(image)
            images.append(image)
            # y_train数据
            mask = cv2.imread(data_mask,cv2.IMREAD_GRAYSCALE)
            mask = cv2.resize(mask, (256,256))
            mask = mask / 255.0
            mask = list(mask)
            masks.append(mask)

        return np.array(images), np.array(masks)


# 读取样本名称，然后根据样本名称去读取数据


train_img = sorted(glob.glob('./trainnsmc/image/*.png'))
train_mask = sorted(glob.glob('./trainnsmc/label/*.png'))
# 数据生成器
training_generator = DataGenerator(train_img, train_mask,batch_size=8)

定义模型，进行训练

model = unet()
#编译模型
from keras_unet_collection import losses
model.compile(optimizer=tf.keras.optimizers.Adam(lr), loss=losses.dice,
              metrics=[ 'acc',losses.dice_coef])

model.fit(training_generator, epochs=50,  max_queue_size=10)

！！！注：本文的代码只展现了关键部分，不是完整代码