🤵♂️ 个人主页:@艾派森的个人主页
✍🏻作者简介:Python学习者
🐋 希望大家多多支持,我们一起进步!😄
如果文章对你有帮助的话,
欢迎评论 💬点赞👍🏻 收藏 📂加关注+
目录
1.项目背景
2.数据集介绍
3.技术工具
4.实验过程
4.1导入数据
4.2数据预处理
4.3特征工程
4.4构建模型
4.5训练模型
4.6模型评估
5.总结
源代码
1.项目背景
在纺织品的自动化生产与质量控制过程中,面料表面常因设备磨损、原料瑕疵或工艺波动产生断经、纬斜、污渍及破洞等各类微小缺陷。这些缺陷如果未能被及时精准地检测出来,不仅会造成原材料的大量浪费,还会直接影响最终产品的质量与品牌信誉。传统的纺织品缺陷检测主要依赖人工目视抽检,这种方式不仅劳动强度大、效率低下,而且极易受到工人主观疲劳和经验差异的影响,难以满足现代工业高质高效的生产要求。随着智能制造与计算机视觉技术的快速发展,利用深度学习模型实现纺织品缺陷的自动精准识别已成为工业界和学术界的研究热点。然而,纺织品背景纹理纷繁复杂,缺陷形态多变且占比极小,对传统卷积神经网络的特征提取能力提出了巨大挑战。为此,本项目基于 Big Transfer (BiT-ResNet50) 强劲的预训练特征表示能力,结合特定数据增强与高性能数据管道,构建了一套高精度的纺织品缺陷智能识别系统,旨在提供一种稳定、高效且具备高泛化能力的工业级缺陷检测解决方案。
2.数据集介绍
本实验数据集来源于Kaggle,为纺织品缺陷检测数据集,利用旋转检测纺织品纹理缺陷,在纺织面料领域,罕见的异常情况可能会发生,从而影响面料的质量。为了避免这种情况的发生,检测缺陷至关重要。本数据集仅用于教学目的。
3.技术工具
Python版本:3.9
代码编辑器:jupyter notebook
4.实验过程
4.1导入数据
在工业级纺织品缺陷检测任务中,精准捕获布料表面的微小瑕疵是保障生产质量的核心。本小节首先导入实验所需的深度学习与数据处理基础库,包括用于构建张量与训练模型的 TensorFlow、加载预训练模型的 TensorFlow Hub,以及用于数据操作与性能指标评估的 H5Py、Pandas、Matplotlib 和 Scikit-Learn 等工具。随后,通过读取训练元数据文件train64.csv,对原始标注字段进行清洗与二值化映射,将复杂的瑕疵类型统一归一化为“损坏(damaged)”与“良好(good)”两类,并将对应的标签数值化为 1 和 0,从而构建出清晰明确的二分类监督学习目标。
# ---------------------------------------------------------------------- # 1. 导入项目所需的第三方依赖库 # ---------------------------------------------------------------------- import os import glob import pickle import tensorflow as tf # 用于将数据集转换为 Tensor 格式以及构建与训练深度学习模型 import tensorflow_hub as hub # 用于下载和调用预训练模型(如 BiT-ResNet50) import numpy as np import h5py # 用于操作和读取 .h5 格式的数据集文件 import pandas as pd import plotly.express as px import matplotlib.pyplot as plt from sklearn.metrics import classification_report # 用于输出模型评估指标报告(如准确率、召回率等) from sklearn.metrics import f1_score # 用于计算模型的 F1 分数 from tqdm import tqdm # 用于在循环或迭代操作中显示动态进度条 # ---------------------------------------------------------------------- # 2. 读取并预处理标注 CSV 文件 # ---------------------------------------------------------------------- # 加载训练集的元数据 CSV 文件 train_df64 = pd.read_csv("/kaggle/input/textiledefectdetection/train64.csv") # 将缺陷类型统一处理:若不为 "good" 则归类为 "damaged"(有缺陷),否则为 "good"(完好) train_df64["indication_type"] = train_df64.indication_type.apply(lambda row: "damaged" if row!="good" else "good") # 将缺陷数值标签映射为二分类数值:非 0 表示有缺陷(标记为 1),0 表示无缺陷(标记为 0) train_df64["indication_value"] = train_df64.indication_value.apply(lambda row: 1 if row!=0 else 0) # 查看数据框的前 5 行,确认格式转换是否正确 train_df64.head()数据读取与标签转换完成后,我们需要对数据集中不同类别(有缺陷与无缺陷)的样本分布情况进行可视化分析。通过计算各类别的占比并绘制柱状图,可以直观地观察数据集是否存在类别不平衡问题,从而指导后续的数据预处理与模型训练策略。
# ---------------------------------------------------------------------- # 3. 计算并可视化不同类别的样本分布比例 # ---------------------------------------------------------------------- # 计算各类别的百分比(保留两位小数),并绘制柱状图 (train_df64.indication_type.value_counts(normalize=True).round(2)*100).to_frame().plot.bar(legend=False, figsize=(8, 6)) # 设置 X 轴和 Y 轴的标签及其字号大小 plt.xlabel('Classes', fontsize=18) plt.ylabel("Count (%)", fontsize=18) # 设置图像标题及其字号大小 plt.title("Distribution of Classes", fontsize=20) # 设置 X 轴刻度标签不旋转(水平显示)及其字号大小 plt.xticks(rotation=0, fontsize=15) plt.show()可以发现数据非常不平衡,需要进行处理!
4.2数据预处理
为了将存放于.h5层级文件中的图像数据高效提取并转化为标准的本地目录结构(以便于 TensorFlow 框架的读取与数据管道构建),我们需要实现一个数据转换类H5ToStorage。该类不仅能通过生成器(Generator)逐条读取图像及其标签信息,还针对训练集进行了数据增强处理(如随机水平/垂直翻转),从而增加无缺陷(good)样本的多样性,平衡数据分布。
# ---------------------------------------------------------------------- # 定义 H5 数据转换与图像提取生成器类 # ---------------------------------------------------------------------- class H5ToStorage: def __init__(self, hdf_path, ds_name="train"): # 保存传入的 HDF5 文件路径 self.path = hdf_path # 初始化类别列表,打开 HDF5 文件获取顶层存储的所有类别名称 self.classes = [] with h5py.File(self.path, 'r') as hf: for class_ in hf: self.classes.append(class_) # 保存数据集类型标识(如 "train" 或 "test") self.name = ds_name # 生成器函数:用于迭代读取 HDF5 文件中的 (图像矩阵, 类别名称, 旋转角度) def generate_img_arr(self): for class_ in self.classes: with h5py.File(self.path, 'r') as hf: for angle in hf[class_]: for img in hf[class_][f"{angle}"]: # 逐条产出图像像素数组、所属类别及对应角度 yield img, class_, angle # 针对训练集生成本地图像目录并进行数据增强 def generate_train_dirs(self): # 创建训练集存放的目标路径结构 path = f"./dataset/{self.name}" os.makedirs(f"{path}/good/") os.makedirs(f"{path}/damaged/") # 注释掉的代码:随机亮度调节层 # random_bright = tf.keras.layers.RandomBrightness(factor=0.05) # 定义 TensorFlow 数据增强层:随机水平与垂直翻转 random_flip = tf.keras.layers.RandomFlip("horizontal_and_vertical") gen = self.generate_img_arr() metadata = {} # 遍历生成器产出的所有图像数据 for i, data in enumerate(gen): img, label, angle = data # 如果是无缺陷(good)样本,则进行 4 次增强扩充,平衡样本数量 if label == "good": for j in range(4): img_path = f"{path}/{label}/{i}_aug{j}.jpeg" # 将图像转为张量并恢复像素范围 [0, 255],应用随机翻转增强 img = random_flip(tf.expand_dims(np.squeeze(img), axis=2)*255., training=True) # 将处理后的图像以灰度图形式保存到本地磁盘 plt.imsave(img_path, np.squeeze(img), cmap="gray") # 如果是有缺陷(damaged)样本,则直接恢复像素范围并保存 else: img_path = f"{path}/damaged/{i}.jpeg" plt.imsave(img_path, np.squeeze(img)*255., cmap="gray") # 记录图像保存路径与角度的映射关系 metadata[img_path] = angle return metadata # 针对测试集生成本地图像目录(不应用数据增强) def generate_test_dirs(self): # 创建测试集存放的目标路径结构 path = f"./dataset/{self.name}" os.makedirs(f"{path}/good/") os.makedirs(f"{path}/damaged/") gen = self.generate_img_arr() metadata = {} # 遍历生成器产出的数据,直接归一化恢复并保持原图保存 for i, data in enumerate(gen): img, label, angle = data if label == "good": img_path = f"{path}/{label}/{i}.jpeg" plt.imsave(img_path, np.squeeze(img)*255., cmap="gray") else: img_path = f"{path}/damaged/{i}.jpeg" plt.imsave(img_path, np.squeeze(img)*255., cmap="gray") # 记录图像保存路径与角度元数据 metadata[img_path] = angle return metadata # 执行入口:根据数据集类型(train / test)调度对应的解析与保存函数 def to_storage(self): if self.name == "train": self.generate_train_dirs() elif self.name == "test": self.generate_test_dirs() # ---------------------------------------------------------------------- # 实例化转换类并分别生成训练集与测试集的本地图像数据 # ---------------------------------------------------------------------- # 实例化训练集生成器对象并执行转换与保存 train_gen = H5ToStorage("/kaggle/input/textiledefectdetection/matchingtDATASET_train_64.h5", "train") train_dict = train_gen.to_storage() # 实例化测试集生成器对象并执行转换与保存 test_gen = H5ToStorage("/kaggle/input/textiledefectdetection/matchingtDATASET_train_64.h5", "test") test_dict = test_gen.to_storage()4.3特征工程
完成本地图像文件的保存后,我们需要借助 TensorFlow 的tf.keras.utils.image_dataset_from_directory接口将硬盘上的图像文件加载为内存中的高性能数据管道(Data Pipeline)。在此阶段,我们设置图像分辨率为 64 x 64,色彩模式为灰度图(grayscale),同时采用prefetch和cache优化内存读取效率,以极大提升后续模型训练时的吞吐速度。
# ---------------------------------------------------------------------- # 1. 定义训练集与测试集本地路径 # ---------------------------------------------------------------------- test_dir = "/kaggle/working/dataset/test" train_dir = "/kaggle/working/dataset/train" # ---------------------------------------------------------------------- # 2. 构建训练集数据管道(Data Pipeline) # ---------------------------------------------------------------------- IMG_SIZE = (64, 64) # 统一将输入图像尺寸调整为 64x64 像素 # 从本地目录批量加载训练集图像,自动根据子文件夹名称(damaged / good)分配整数标签 train_data = tf.keras.utils.image_dataset_from_directory(directory=train_dir, image_size=IMG_SIZE, color_mode="grayscale", label_mode="int", # 标签编码形式为整数数字类型 [0, 1] batch_size=64) # 设定批次大小(Batch Size)为 64 # 获取并打印训练集中的类别名称映射 class_names = train_data.class_names print(f"The class_names are as follows: {class_names}") # 启用数据管道性能优化:利用预取(prefetch)和缓存(cache)加速 GPU 训练过程中的数据加载 train_data = train_data.prefetch(buffer_size=tf.data.AUTOTUNE).\ cache() # ---------------------------------------------------------------------- # 3. 构建测试集数据管道 # ---------------------------------------------------------------------- # 从本地目录加载测试集图像,设置为不打乱顺序(shuffle=False),确保预测结果与真实标签对应 test_data = tf.keras.utils.image_dataset_from_directory(directory=test_dir, image_size=IMG_SIZE, color_mode="grayscale", shuffle=False, batch_size=64, label_mode="int") # 针对测试集数据管道开启 prefetch 与 cache 性能优化 test_data = test_data.prefetch(buffer_size=tf.data.AUTOTUNE).\ cache()经过数据预处理阶段的数据增强(对无缺陷图像进行翻转扩充)后,我们需要再次统计并绘制当前训练集中“有缺陷(damaged)”与“无缺陷(good)”图像的数量占比,以验证增强后的数据集类别比例是否得到显著改善。
# ---------------------------------------------------------------------- # 4. 重新统计经过增强处理后的训练集类别数量与占比 # ---------------------------------------------------------------------- # 统计 damaged 子目录下扩充后的图像总数 damaged_counts = 0 for _ in os.listdir(f"{train_dir}/damaged"): damaged_counts += 1 # 统计 good 子目录下扩充后的图像总数 good_counts = 0 for _ in os.listdir(f"{train_dir}/good"): good_counts += 1 # 计算训练集图像总数 total = damaged_counts + good_counts # 构建包含各类占比数据的 Pandas Series 对象,并四舍五入保留两位小数 new_train_df64 = pd.Series({"damaged": damaged_counts/total, "good": good_counts/total}).round(2) # ---------------------------------------------------------------------- # 5. 绘制增强后的类别分布百分比柱状图 # ---------------------------------------------------------------------- (new_train_df64*100).to_frame().plot.bar(legend=False, figsize=(8, 6)) plt.xlabel('Classes', fontsize=18) plt.ylabel("Count (%)", fontsize=18) plt.title("Distribution of Classes", fontsize=20) plt.xticks(rotation=0, fontsize=15) plt.show()为了直观检查图像数据管道加载的效果以及图像像素是否保持正常,我们编写一个可视化函数explore_data。该函数通过随机打乱(shuffle)并抽取部分批次图像,以多行多列网格(Grid)的形式随机展示训练集中的图像及其对应的分类标签。
# ---------------------------------------------------------------------- # 6. 定义并执行训练集图像可视化数据探索函数 # ---------------------------------------------------------------------- def explore_data(n_rows, n_cols, figsize): count = 0 # 创建指定行数和列数的子图画布 fig, axs = plt.subplots(n_rows, n_cols, figsize=figsize) fig.suptitle("Training Data Exploration", fontsize=18, y=0.95) # 从打乱后的数据集中随机抽取 1 个 batch 并解包(unbatch),与子图网格一一对应绑定 for data, ax in zip(train_data.shuffle(10000).take(1).unbatch(), axs.ravel()): img, label = data # 关闭坐标轴显示 ax.axis("off") # 将子图标题设置为具体的类别名称(damaged 或 good) ax.set_title(f"{class_names[label]}") # 以灰度色彩映射显示单通道图像 ax.imshow(np.squeeze(img), cmap="gray") count += 1 # 当绘制数量达到预设的子图网格总量(n_rows * n_cols)时停止循环 if count == n_rows * n_cols: break; # 调用探索函数,绘制 2 行 6 列共 12 张图像样本进行可视化检验 explore_data(n_rows=2, n_cols=6, figsize=(15, 6))4.4构建模型
在模型构建阶段,我们借助迁移学习(Transfer Learning)技术,调用 TensorFlow Hub 上预训练的BiT-ResNet50 (Big Transfer ResNet50)作为基础特征提取器。由于 BiT-ResNet50 模型的标准输入需要 3 通道(RGB)图像,而我们的数据集为 1 通道灰度图,因此我们在模型最前端巧用 1 x 1 的卷积层(Conv2D)将通道数从 1 扩展至 3。随后,接上预训练特征提取层(设置为不可训练trainable=False以冻结权重),最后叠加全连接层与 Sigmoid 激活函数,构建出一个高效的二分类识别网络。
# ---------------------------------------------------------------------- # 1. 从 TensorFlow Hub 加载预训练的 BiT-ResNet50 特征提取器 # ---------------------------------------------------------------------- import tensorflow_hub as hub # 加载基于 ImageNet-21k 预训练的 BiT-ResNet50 模型层,并将权重冻结(trainable=False) module = hub.KerasLayer("https://tfhub.dev/google/experts/bit/r50x1/in21k/angiosperm/1", trainable=False, name='feature_extraction_layer', input_shape=IMG_SIZE+(3,)) # ---------------------------------------------------------------------- # 2. 构建 Sequential 顺序分类模型 # ---------------------------------------------------------------------- model = tf.keras.Sequential( [ # 使用 1x1 卷积将单通道灰度图 (64, 64, 1) 升维映射为 3 通道 (64, 64, 3),以适配 BiT-ResNet50 的输入格式需求 tf.keras.layers.Conv2D(3, 1, padding="same", activation="relu"), # 预训练特征提取层(提取图像的高阶抽象特征) module, # 隐藏全连接层:256 个神经元,采用 ReLU 激活函数 tf.keras.layers.Dense(256, activation="relu", name="dense_layer"), # 输出层:单个神经元,采用 Sigmoid 激活函数用于二分类(概率输出 0~1) tf.keras.layers.Dense(1, activation='sigmoid', name='output_layer') ] ) # ---------------------------------------------------------------------- # 3. 显式指定输入形状构建模型并打印模型结构摘要 # ---------------------------------------------------------------------- # 传入单通道图像张量形状 (Batch_Size, Height, Width, Channels) model.build((None, 64, 64, 1)) # 查看网络各层参数量与结构细节 model.summary()构建好网络拓扑结构后,我们可以使用 Keras 提供的plot_model工具将模型架构绘制成直观的流程图。这不仅方便检查输入/输出张量在各个网络层之间的形状变化(Shapes),还能在博客中更清晰地向读者展示模型的整体层次结构。
# ---------------------------------------------------------------------- # 4. 可视化绘制模型网络结构拓扑图 # ---------------------------------------------------------------------- tf.keras.utils.plot_model( model, # to_file='model.png', # 如需保存图片到本地,取消此行注释即可 show_shapes=True, # 显示各层的输入输出张量形状(Shape) show_dtype=False, # 不显示数据类型 show_layer_names=True, # 显示各层的名称 rankdir='TB', # 图像排布方向:从上到下(Top to Bottom) expand_nested=False, # 不展开嵌套模型 dpi=96, # 图像输出分辨率 layer_range=None, # 绘制所有层 )4.5训练模型
在完成模型结构的搭建后,我们需要对其进行编译(Compile)。由于本项目是二分类任务,损失函数选择二元交叉熵(BinaryCrossentropy),优化器采用Adam算法,初始学习率设为 0.01,并选用准确率(accuracy)作为评估指标。在正式开始训练(fit)时,我们将训练集数据输入网络,设定训练轮数为 10,并传入测试集作为验证集来实时监控模型的泛化性能。同时引入回调函数(Callbacks)以实现动态调整学习率与提前终止训练,防止模型过拟合。
# ---------------------------------------------------------------------- # 1. 配置损失函数、优化器与评估指标(编译模型) # ---------------------------------------------------------------------- # 定义二分类交叉熵损失函数(因为输出层已使用 Sigmoid 激活函数,所以 from_logits=False) binary_categorical_loss = tf.keras.losses.BinaryCrossentropy(from_logits=False) # 定义 Adam 优化器,设置初始学习率为 0.01 optimizer = tf.keras.optimizers.Adam(learning_rate=0.01) # 编译模型,组装损失函数、优化器和准确率评估指标 model.compile( loss=binary_categorical_loss, optimizer=optimizer, metrics=["accuracy"] ) # ---------------------------------------------------------------------- # 2. 设置训练超参数并启动模型训练 # ---------------------------------------------------------------------- EPOCHS = 10 # 设定总训练轮数为 10 STEPS_PER_EPOCH = len(train_data) # 计算每个 Epoch 的迭代步数(等于训练集批次总量) # 执行模型训练,并将训练过程中的 Loss 和 Accuracy 历史指标保存至 training_history_1 training_history_1 = model.fit(train_data, epochs=EPOCHS, steps_per_epoch=STEPS_PER_EPOCH, validation_data=test_data.shuffle(10000), # 使用打乱后的测试集作为验证集 validation_steps=int(0.5 * len(test_data)), # 每次验证时仅使用 50% 的测试数据进行快速评估 callbacks=[reduce_lr, early_stopping] # 加入回调函数以动态调参和防止过拟合 )4.6模型评估
在模型训练完成后,我们需要在测试集上对其进行全面评估。评估过程分为三个阶段:首先在整个测试集上计算最终的损失值(Loss)与准确率(Accuracy);其次通过绘制训练与验证曲线直观观察模型的收敛轨迹;最后获取测试集的真实标签与模型预测概率,生成详细的分类评估报告(包含精确率 Precision、召回率 Recall 以及 F1 分数)。
首先,调用model.evaluate函数在测试集上进行整体性能评估:
# ---------------------------------------------------------------------- # 1. 在完整的测试数据集上评估模型性能,输出 Loss 和 Accuracy # ---------------------------------------------------------------------- model.evaluate(test_data)其次,传入训练历史记录training_history_1,绘制训练集与验证集在各个 Epoch 下的损失曲线和准确率变化曲线:
# ---------------------------------------------------------------------- # 2. 绘制训练与验证过程中的评估曲线(损失曲线与准确率曲线) # ---------------------------------------------------------------------- plot_evaluation_curves(training_history_1, 2)最后,获取测试集的真实标签y_labels_1与预测概率y_preds_1,将连续的预测概率通过四舍五入(np.round)转化为 0/1 二分类预测值,并打印分类性能报告:
# ---------------------------------------------------------------------- # 3. 提取测试集的真实标签与模型预测值,生成详细分类评估报告 # ---------------------------------------------------------------------- # 调用自定义函数提取测试集中的真实标签 (y_labels_1) 和预测概率 (y_preds_1) y_labels_1, y_preds_1 = get_labels(test_data, model) # 打印分类报告表头 print("\n\n############### Classification Report ###############\n\n") # 对预测概率进行四舍五入二值化(>=0.5 判定为 1,<0.5 判定为 0),并计算各项评估指标 print(classification_report(y_labels_1, np.round(y_preds_1)))5.总结
本项目针对纺织面料生产中偶发性缺陷难以高效识别的痛点,基于 Kaggle 纺织品缺陷检测数据集构建并实现了一套基于 BiT-ResNet50 迁移学习的智能图像识别方案。通过对 HDF5 格式图像数据的离线解析与数据增强,有效克服了原始样本类别不平衡的问题,并结合 $1 \times 1$ 卷积降维映射与 TensorFlow 高性能数据管道,完成了高效的模型训练与验证。实验结果表明,该模型在包含 72,000 个测试样本的验证集上展现出了极其卓越的泛化能力与鲁棒性,测试集 Loss 低至 0.0081,整体准确率达到 99.74%,且各类别的 Precision、Recall 和 F1-score 均保持在 0.99 以上的高水准,充分证实了 BiT-ResNet50 架构在纺织品微小缺陷检测任务中的实用价值与高精度表现。
源代码
# import the libraries import os import glob import pickle import tensorflow as tf # used for creating tensors out of the dataset and training the model import tensorflow_hub as hub # downloading and using pre-trained models import numpy as np import h5py # to manipulate the dataset as it's in `.h5` format import pandas as pd import plotly.express as px import matplotlib.pyplot as plt from sklearn.metrics import classification_report # essential for model's evaluation from sklearn.metrics import f1_score from tqdm import tqdm # getting a progress bar for an iterable object # load the `train64.csv` file train_df64 = pd.read_csv("/kaggle/input/textiledefectdetection/train64.csv") train_df64["indication_type"] = train_df64.indication_type.apply(lambda row: "damaged" if row!="good" else "good") train_df64["indication_value"] = train_df64.indication_value.apply(lambda row: 1 if row!=0 else 0) train_df64.head() # plot the distribution of classes (train_df64.indication_type.value_counts(normalize=True).round(2)*100).to_frame().plot.bar(legend=False, figsize=(8, 6)) plt.xlabel('Classes', fontsize=18) plt.ylabel("Count (%)", fontsize=18) plt.title("Distribution of Classes", fontsize=20) plt.xticks(rotation=0, fontsize=15); # create an object that will take the dataset and produce the dataset in a format required for tensorflow dataset's API class H5ToStorage: def __init__(self, hdf_path, ds_name="train"): self.path = hdf_path self.classes = [] with h5py.File(self.path, 'r') as hf: for class_ in hf: self.classes.append(class_) self.name = ds_name # a generator to load the (img, class, angle) def generate_img_arr(self): for class_ in self.classes: with h5py.File(self.path, 'r') as hf: for angle in hf[class_]: for img in hf[class_][f"{angle}"]: yield img, class_, angle # utilize the generator to create new images and load it back to Storage def generate_train_dirs(self): # create the dataset's directories path = f"./dataset/{self.name}" os.makedirs(f"{path}/good/") os.makedirs(f"{path}/damaged/") # random_bright = tf.keras.layers.RandomBrightness(factor=0.05) random_flip = tf.keras.layers.RandomFlip("horizontal_and_vertical") gen = self.generate_img_arr() metadata = {} for i, data in enumerate(gen): img, label, angle = data if label == "good": for j in range(4): img_path = f"{path}/{label}/{i}_aug{j}.jpeg" img = random_flip(tf.expand_dims(np.squeeze(img), axis=2)*255., training=True) plt.imsave(img_path, np.squeeze(img), cmap="gray") else: img_path = f"{path}/damaged/{i}.jpeg" plt.imsave(img_path, np.squeeze(img)*255., cmap="gray") metadata[img_path] = angle return metadata def generate_test_dirs(self): # create the dataset's directories path = f"./dataset/{self.name}" os.makedirs(f"{path}/good/") os.makedirs(f"{path}/damaged/") gen = self.generate_img_arr() metadata = {} for i, data in enumerate(gen): img, label, angle = data if label == "good": img_path = f"{path}/{label}/{i}.jpeg" plt.imsave(img_path, np.squeeze(img)*255., cmap="gray") else: img_path = f"{path}/damaged/{i}.jpeg" plt.imsave(img_path, np.squeeze(img)*255., cmap="gray") metadata[img_path] = angle return metadata def to_storage(self): if self.name == "train": self.generate_train_dirs() elif self.name == "test": self.generate_test_dirs() # generate train data train_gen = H5ToStorage("/kaggle/input/textiledefectdetection/matchingtDATASET_train_64.h5", "train") train_dict = train_gen.to_storage() # generate test data test_gen = H5ToStorage("/kaggle/input/textiledefectdetection/matchingtDATASET_train_64.h5", "test") test_dict = test_gen.to_storage() # train data & test data paths test_dir = "/kaggle/working/dataset/test" train_dir = "/kaggle/working/dataset/train" # create train data pipeline IMG_SIZE = (64, 64) # define image size train_data = tf.keras.utils.image_dataset_from_directory(directory=train_dir, image_size=IMG_SIZE, color_mode="grayscale", label_mode="int", # label type is encoding of the form `[1, 2, 3, ...]` batch_size=64) # batch_size is 32 by default, but 128 will be used for this use-case # classes in train_data class_names = train_data.class_names print(f"The class_names are as follows: {class_names}") train_data = train_data.prefetch(buffer_size=tf.data.AUTOTUNE).\ cache() # create test data pipeline test_data = tf.keras.utils.image_dataset_from_directory(directory=test_dir, image_size=IMG_SIZE, color_mode="grayscale", shuffle=False, batch_size=64, label_mode="int") test_data = test_data.prefetch(buffer_size=tf.data.AUTOTUNE).\ cache() # check the proportions of classes damaged_counts = 0 for _ in os.listdir(f"{train_dir}/damaged"): damaged_counts += 1 good_counts = 0 for _ in os.listdir(f"{train_dir}/good"): good_counts += 1 total = damaged_counts + good_counts new_train_df64 = pd.Series({"damaged": damaged_counts/total, "good": good_counts/total}).round(2) # check the classes distribution (new_train_df64*100).to_frame().plot.bar(legend=False, figsize=(8, 6)) plt.xlabel('Classes', fontsize=18) plt.ylabel("Count (%)", fontsize=18) plt.title("Distribution of Classes", fontsize=20) plt.xticks(rotation=0, fontsize=15) plt.show() # explore the train data def explore_data(n_rows, n_cols, figsize): count = 0 fig, axs = plt.subplots(n_rows, n_cols, figsize=figsize) # plot 20 figures fig.suptitle("Training Data Exploration", fontsize=18, y=0.95) for data, ax in zip(train_data.shuffle(10000).take(1).unbatch(), axs.ravel()): img, label = data ax.axis("off") ax.set_title(f"{class_names[label]}") ax.imshow(np.squeeze(img), cmap="gray") count+=1 if count == n_rows*n_cols: break; explore_data(n_rows=2, n_cols=6, figsize=(15, 6)) import tensorflow_hub as hub # feature extractor trained on angiosperms module = hub.KerasLayer("https://tfhub.dev/google/experts/bit/r50x1/in21k/angiosperm/1", trainable=False, name='feature_extraction_layer', input_shape=IMG_SIZE+(3,)) # build the sequential model model = tf.keras.Sequential( [ tf.keras.layers.Conv2D(3, 1, padding="same", activation="relu"), # the following feature extraction layer won't accept the data of shape `(None, 64, 64, 1)` # to solve this issue, it is required that the channels are expanded form `1` to `3` by using 1x1 convolution module, # feature extractor layer tf.keras.layers.Dense(256, activation="relu", name="dense_layer"), tf.keras.layers.Dense(1, activation='sigmoid', name='output_layer') # making predictions ] ) # build the model and check summary model.build((None, 64, 64, 1)) model.summary() # plot the model tf.keras.utils.plot_model( model, # to_file='model.png', show_shapes=True, show_dtype=False, show_layer_names=True, rankdir='TB', expand_nested=False, dpi=96, layer_range=None, ) # Define the loss function binary_categorical_loss = tf.keras.losses.BinaryCrossentropy(from_logits=False) # Define the optimizer optimizer = tf.keras.optimizers.Adam(learning_rate=0.01) # Compile the model model.compile( loss=binary_categorical_loss, optimizer=optimizer, metrics=["accuracy"] ) # train the model EPOCHS = 10 STEPS_PER_EPOCH = len(train_data) training_history_1 = model.fit(train_data, epochs=EPOCHS, steps_per_epoch=STEPS_PER_EPOCH, validation_data=test_data.shuffle(10000), validation_steps=int(0.5 * len(test_data)), callbacks=[reduce_lr, early_stopping] ) model.evaluate(test_data) plot_evaluation_curves(training_history_1, 2) # create classification_report y_labels_1, y_preds_1 = get_labels(test_data, model) print("\n\n############## Classification Report ###############\n\n") print(classification_report(y_labels_1, np.round(y_preds_1)))资料获取,更多粉丝福利,关注下方公众号获取