Flutter机器学习推理插件pytorch_lite的使用

pytorch_lite 是一个Flutter包，旨在帮助运行PyTorch Lite模型进行分类和目标检测（包括YOLOV5和YOLOV8）。以下是详细的使用指南，包括如何准备模型、安装依赖以及编写完整的示例代码。

准备模型

分类模型

import torch
from torch.utils.mobile_optimizer import optimize_for_mobile

model = torch.load('model_scripted.pt', map_location="cpu")
model.eval()
example = torch.rand(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example)
optimized_traced_model = optimize_for_mobile(traced_script_module)
optimized_traced_model._save_for_lite_interpreter("model.pt")

目标检测模型 (YOLOV5)

!python export.py --weights "the weights of your model" --include torchscript --img 640 --optimize

例如：

!python export.py --weights yolov5s.pt --include torchscript --img 640 --optimize

目标检测模型 (YOLOV8)

!yolo mode=export model="your model" format=torchscript optimize

例如：

!yolo mode=export model=yolov8s.pt format=torchscript optimize

安装

在 pubspec.yaml 文件中添加 pytorch_lite 作为依赖项：

dependencies:
  pytorch_lite: ^latest_version

创建一个 assets 文件夹，并将你的 PyTorch 模型和标签文件放入其中。修改 pubspec.yaml 文件以包含这些资源：

assets:
  - assets/models/model_classification.pt
  - assets/labels_classification.txt
  - assets/models/model_objectDetection.torchscript
  - assets/labels_objectDetection.txt

运行 flutter pub get。

对于发布版本，在 android/app/build.gradle 文件中添加以下配置：

buildTypes {
    release {
        shrinkResources false
        minifyEnabled false
        signingConfig signingConfigs.debug
    }
}

使用

导入库

import 'package:pytorch_lite/pytorch_lite.dart';

加载模型

分类模型

ClassificationModel classificationModel = await PytorchLite.loadClassificationModel(
  "assets/models/model_classification.pt", 
  224, 224,
  labelPath: "assets/labels/label_classification_imageNet.txt"
);

目标检测模型

ModelObjectDetection objectModel = await PytorchLite.loadObjectDetectionModel(
  "assets/models/yolov5s.torchscript", 
  80, 640, 640,
  labelPath: "assets/labels/labels_objectDetection_Coco.txt",
  objectDetectionModelType: ObjectDetectionModelType.yolov5
);

获取分类预测结果

从图片获取预测结果

String imagePrediction = await classificationModel.getImagePrediction(
  await File(image.path).readAsBytes()
);

从摄像头图像获取预测结果

String imagePrediction = await _objectModel.getCameraImagePrediction(
  cameraImage,
  rotation, // 检查示例中的旋转值
);

获取原始输出层

List<double>? predictionList = await _imageModel!.getImagePredictionList(
  await File(image.path).readAsBytes(),
);

从摄像头图像获取原始输出层

List<double>? predictionList = await _imageModel!.getCameraImagePredictionList(
  cameraImage,
  rotation, // 检查示例中的旋转值
);

获取概率（如果模型未使用softmax）

List<double>? predictionListProbabilities = await _imageModel!.getImagePredictionListProbabilities(
  await File(image.path).readAsBytes(),
);

从摄像头图像获取概率

List<double>? predictionListProbabilities = await _imageModel!.getCameraPredictionListProbabilities(
  cameraImage,
  rotation, // 检查示例中的旋转值
);

获取目标检测预测结果

从图片获取预测结果

List<ResultObjectDetection> objDetect = await _objectModel.getImagePrediction(
  await File(image.path).readAsBytes(),
  minimumScore: 0.1, 
  iOUThreshold: 0.3
);

从摄像头图像获取预测结果

List<ResultObjectDetection> objDetect = await _objectModel.getCameraImagePrediction(
  cameraImage,
  rotation, // 检查示例中的旋转值
  minimumScore: 0.1, 
  iOUThreshold: 0.3
);

在图像上绘制边界框

objectModel.renderBoxesOnImage(_image!, objDetect);

自定义均值和标准差的图像预测

final mean = [0.5, 0.5, 0.5];
final std = [0.5, 0.5, 0.5];
String prediction = await classificationModel.getImagePrediction(
  image, 
  mean: mean, 
  std: std
);

示例代码

以下是一个完整的示例代码，展示了如何在Flutter应用中使用 pytorch_lite 插件：

import 'package:flutter/material.dart';
import 'package:pytorch_lite/pytorch_lite.dart';

void main() async {
  runApp(const ChooseDemo());
}

class ChooseDemo extends StatefulWidget {
  const ChooseDemo({Key? key}) : super(key: key);

  @override
  State<ChooseDemo> createState() => _ChooseDemoState();
}

class _ChooseDemoState extends State<ChooseDemo> {
  late ClassificationModel classificationModel;
  late ModelObjectDetection objectModel;

  @override
  void initState() {
    super.initState();
    loadModels();
  }

  Future<void> loadModels() async {
    classificationModel = await PytorchLite.loadClassificationModel(
      "assets/models/model_classification.pt", 
      224, 224,
      labelPath: "assets/labels/label_classification_imageNet.txt"
    );

    objectModel = await PytorchLite.loadObjectDetectionModel(
      "assets/models/yolov5s.torchscript", 
      80, 640, 640,
      labelPath: "assets/labels/labels_objectDetection_Coco.txt",
      objectDetectionModelType: ObjectDetectionModelType.yolov5
    );
  }

  @override
  Widget build(BuildContext context) {
    return MaterialApp(
      home: Scaffold(
        appBar: AppBar(
          title: const Text('Pytorch Mobile Example'),
        ),
        body: Builder(builder: (context) {
          return Center(
            child: Column(
              children: [
                TextButton(
                  onPressed: () async {
                    // 这里可以添加逻辑来处理模型推理
                    String prediction = await classificationModel.getImagePrediction(
                      await File(image.path).readAsBytes()
                    );
                    print(prediction);
                  },
                  style: TextButton.styleFrom(
                    backgroundColor: Colors.blue,
                  ),
                  child: const Text(
                    "Run Classification Model",
                    style: TextStyle(
                      color: Colors.white,
                    ),
                  ),
                ),
                TextButton(
                  onPressed: () async {
                    // 这里可以添加逻辑来处理模型推理
                    List<ResultObjectDetection> objDetect = await objectModel.getImagePrediction(
                      await File(image.path).readAsBytes(),
                      minimumScore: 0.1, 
                      iOUThreshold: 0.3
                    );
                    print(objDetect);
                  },
                  style: TextButton.styleFrom(
                    backgroundColor: Colors.blue,
                  ),
                  child: const Text(
                    "Run Object Detection Model",
                    style: TextStyle(
                      color: Colors.white,
                    ),
                  ),
                ),
              ],
            ),
          );
        }),
      ),
    );
  }
}

通过以上步骤，你可以轻松地在Flutter应用中集成并使用 pytorch_lite 插件进行机器学习推理。希望这些信息对你有所帮助！