使用 ML.NET 实现图像分类：从入门到实践

ImageClassification/
├── Program.cs
├── assets/           # 存放训练图片
│   ├── CD/          # 有裂缝的图片
│   └── UD/          # 无裂缝的图片  
└── workspace/        # 存放模型文件

// 原始图像数据类
public class ImageData
{
    public string ImagePath { get; set; }
    public string Label { get; set; }
}

// 模型输入类
public class ModelInput
{
    public byte[] Image { get; set; }
    public UInt32 LabelAsKey { get; set; }
    public string ImagePath { get; set; }
    public string Label { get; set; }
}

// 模型输出类
public class ModelOutput
{
    public string ImagePath { get; set; }
    public string Label { get; set; }
    public string PredictedLabel { get; set; }
}

private static IEnumerable<ImageData> LoadImagesFromDirectory(string folder, bool useFolderNameAsLabel = true)
{
    var files = Directory.GetFiles(folder, "*", searchOption: SearchOption.AllDirectories);

    foreach (var file in files)
    {
        if ((Path.GetExtension(file) != ".jpg") && (Path.GetExtension(file) != ".png"))
            continue;

        var label = Path.GetFileName(file);
        if (useFolderNameAsLabel)
            label = Directory.GetParent(file).Name;
        else
        {
            for (int index = 0; index < label.Length; index++)
            {
                if (!char.IsLetter(label[index]))
                {
                    label = label.Substring(0, index);
                    break;
                }
            }
        }

        yield return new ImageData()
        {
            ImagePath = file,
            Label = label
        };
    }
}

class Program
{
    static void Main(string[] args)
    {
        // 初始化ML.NET环境
        MLContext mlContext = new MLContext();

        // 设置路径
        var projectDirectory = Path.GetFullPath(Path.Combine(AppContext.BaseDirectory, "../../../"));
        var workspaceRelativePath = Path.Combine(projectDirectory, "workspace");
        var assetsRelativePath = Path.Combine(projectDirectory, "assets");

        // 加载数据
        IEnumerable<ImageData> images = LoadImagesFromDirectory(folder: assetsRelativePath, useFolderNameAsLabel: true);
        IDataView imageData = mlContext.Data.LoadFromEnumerable(images);
        IDataView shuffledData = mlContext.Data.ShuffleRows(imageData);

        // 数据预处理
        var preprocessingPipeline = mlContext.Transforms.Conversion.MapValueToKey(
            inputColumnName: "Label",
            outputColumnName: "LabelAsKey")
        .Append(mlContext.Transforms.LoadRawImageBytes(
            outputColumnName: "Image",
            imageFolder: assetsRelativePath,
            inputColumnName: "ImagePath"));

        IDataView preProcessedData = preprocessingPipeline
            .Fit(shuffledData)
            .Transform(shuffledData);

        // 数据集分割
        var trainSplit = mlContext.Data.TrainTestSplit(data: preProcessedData, testFraction: 0.3);
        var validationTestSplit = mlContext.Data.TrainTestSplit(trainSplit.TestSet);

        // 配置训练选项
        var classifierOptions = new ImageClassificationTrainer.Options()
        {
            FeatureColumnName = "Image",
            LabelColumnName = "LabelAsKey",
            ValidationSet = validationTestSplit.TrainSet,
            Arch = ImageClassificationTrainer.Architecture.ResnetV2101,
            MetricsCallback = (metrics) => Console.WriteLine(metrics),
            TestOnTrainSet = false,
            ReuseTrainSetBottleneckCachedValues = true,
            ReuseValidationSetBottleneckCachedValues = true,
            WorkspacePath = workspaceRelativePath
        };

        // 定义训练管道
        var trainingPipeline = mlContext.MulticlassClassification.Trainers
            .ImageClassification(classifierOptions)
            .Append(mlContext.Transforms.Conversion.MapKeyToValue("PredictedLabel"));

        // 训练模型
        Console.WriteLine("*** 开始训练模型 ***");
        ITransformer trainedModel = trainingPipeline.Fit(trainSplit.TrainSet);

        // 进行预测
        ClassifySingleImage(mlContext, validationTestSplit.TestSet, trainedModel);
        ClassifyImages(mlContext, validationTestSplit.TestSet, trainedModel);
    }
}

mlContext.Transforms.Conversion.MapValueToKey(
    inputColumnName: "Label",
    outputColumnName: "LabelAsKey")

- 将字符串类型的标签("Label")转换为数值类型("LabelAsKey")
- 例如："CD" -> 0, "UD" -> 1
- 这是必需的，因为机器学习模型需要数值形式的标签
- 输入是 ImageData 类中的 Label 属性
- 输出存储在 ModelInput 类的 LabelAsKey 属性中

mlContext.Transforms.LoadRawImageBytes(
    outputColumnName: "Image",
    imageFolder: assetsRelativePath,
    inputColumnName: "ImagePath")

- 将图像文件转换为字节数组格式
- `outputColumnName: "Image"`: 输出到 ModelInput 类的 Image 属性
- `imageFolder: assetsRelativePath`: 指定图像文件所在的根目录
- `inputColumnName: "ImagePath"`: 使用 ImageData 类中的 ImagePath 属性

mlContext.MulticlassClassification.Trainers.ImageClassification(classifierOptions)

- 使用多分类分类器进行图像分类
- 基于之前定义的 classifierOptions 配置
- 使用迁移学习方法，基于预训练的 ResNet 模型
- 主要功能：
    - 提取图像特征
    - 训练分类器
    - 生成预测模型

.Append(mlContext.Transforms.Conversion.MapKeyToValue("PredictedLabel"))

- 将模型输出的数值预测结果转换回原始标签
- 与前面的 MapValueToKey 操作相反
- 例如：将 0 转回 "CD"，1 转回 "UD"
- 确保最终输出是人类可读的标签

private static void ClassifySingleImage(MLContext mlContext, IDataView data, ITransformer trainedModel)
{
    var predictionEngine = mlContext.Model.CreatePredictionEngine<ModelInput, ModelOutput>(trainedModel);
    var image = mlContext.Data.CreateEnumerable<ModelInput>(data, reuseRowObject: true).First();
    var prediction = predictionEngine.Predict(image);

    Console.WriteLine($"单张图片分类结果:");
    Console.WriteLine($"图片: {Path.GetFileName(prediction.ImagePath)}");
    Console.WriteLine($"实际类别: {prediction.Label}");
    Console.WriteLine($"预测类别: {prediction.PredictedLabel}");
}

private static void ClassifyImages(MLContext mlContext, IDataView data, ITransformer trainedModel)
{
    IDataView predictionData = trainedModel.Transform(data);
    var predictions = mlContext.Data.CreateEnumerable<ModelOutput>(predictionData, reuseRowObject: true)
        .Take(10);

    Console.WriteLine("\n批量图片分类结果:");
    foreach (var prediction in predictions)
    {
        Console.WriteLine($"图片: {Path.GetFileName(prediction.ImagePath)}");
        Console.WriteLine($"实际类别: {prediction.Label}");
        Console.WriteLine($"预测类别: {prediction.PredictedLabel}\n");
    }
}