MicroTiClassifier

MicroTiClassifier 是一个基于深度学习的图像分类实验项目。

基于预训练的 VGG16 模型，使用 5 类共 825 张原始图片进行迁移学习，实现对钛合金微观组织图像的分类。

项目结构

模块	功能	详细
config	配置文件	文件路径、训练参数
data	数据集	无
image	示例图片	无
log	数据处理日志	运行时生成
output	模型训练输出	运行时生成
scripts	可视化分析脚本	网格图、训练日志、混淆矩阵、PCA 降维、PR 曲线、分类报告、ROC 曲线、数据集统计、t-SNE 降维、UMAP 降维
source	模型相关脚本	模型导出、类激活热力图、模型预测、语义特征提取、模型训练
utils	数据处理脚本	原始数据集分析、数据增强、图像文件筛选、数据集划分、数据集统计分析

实验环境

OS	Python	PyTorch	CUDA
Windows 11 24h2	3.11	2.3.1	-
Ubuntu 22.04 LTS	3.11	2.3.1	12.1

工作流程

flowchart TD
    subgraph SP1[数据预处理]
        direction LR

        A[原始数据集]
        B[analyze.py<br/>分析原始数据集]
        C[filter.py<br/>筛选图像文件]
        D[人工分类形成分类数据集]
        E[split.py<br/>划分训练集和验证集]

        A --> B --> C --> D --> E
    end

    subgraph SP2[数据增强]
        direction LR

        F[augment.py<br/>数据增强]
        F1[utils/stats.py<br/>数据增强前后文件数量统计]
        F11[scripts/stats.py<br/>统计分析可视化]

        F --> F1 --> F11
    end

    subgraph SP3[模型训练]
        direction LR

        G[train.py<br/>训练模型]
        G1[log.py<br/>训练日志可视化]
        G2[export.py<br/>导出为 onnx 格式]

        G --> G1 --> G2
    end

    subgraph SP4[性能评估]
        direction LR

        H[predict.py<br/>在验证集上预测]
        H1[matrix.py, PR.py, ROC.py<br/>混淆矩阵，PR，ROC]
        H2[semantic.py<br/>语义特征提取]
        H21[PCA.py, t-SNE.py, UMAP.py<br/>语义特征降维可视化]
        H3[grad-cam.py<br/>类激活热力图]

        H --> H1
        H --> H2 --> H21
        H --> H3
    end

    SP1 --> SP2 --> SP3 --> SP4

Loading

快速开始

1. 克隆

git clone https://github.com/huang2fire/MicroTiClassifier.git

2. 环境

推荐使用 miniforge 进行环境配置。

conda create -n MicroTiClassifier python=3.11
conda activate MicroTiClassifier
conda install pytorch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 pytorch-cuda=12.1 -c pytorch -c nvidia
conda install grad-cam matplotlib pandas scikit-learn seaborn tqdm transformers umap-learn

3. 数据集

暂不开放。

4. 模型

详见 releases

许可证

本项目采用 MIT 许可证，详见 LICENSE 文件。