313 lines
11 KiB
Dart
313 lines
11 KiB
Dart
import 'dart:typed_data';
|
||
|
||
import 'package:tflite_flutter/tflite_flutter.dart';
|
||
|
||
import 'detection_result.dart';
|
||
import 'nms.dart';
|
||
|
||
/// YOLOv8n 端侧推理实现（对应 Kotlin TFLiteDetector）。
|
||
/// 模型输出布局（ultralytics litert 导出）：[1, 4 + nc, anchors]，
|
||
/// cx/cy/w/h 已归一化，类别得分已过 sigmoid；按 out[dim][anchor] 索引。
|
||
/// 输入为 NCHW [1, 3, 704, 704]（litert 导出保留 torch 布局）。
|
||
class TfliteDetector {
|
||
static const int inputSize = 704;
|
||
// 野鸡数据置信度普遍偏低（0.1~0.2 量级），保留低分池供运动检测提升
|
||
static const double minScore = 0.10;
|
||
static const double iouThreshold = 0.45;
|
||
static const int maxDetections = 20;
|
||
static const String modelAsset = 'assets/model.tflite';
|
||
static const String labelsAsset = 'assets/labels.txt';
|
||
|
||
final Interpreter _interpreter;
|
||
final List<String> _labels;
|
||
final int _numClasses;
|
||
final int _numAnchors;
|
||
|
||
final Float32List _input =
|
||
Float32List(1 * inputSize * inputSize * 3);
|
||
|
||
/// 输出按模型形状 [1, 4+nc, anchors] 的嵌套 List 组织，
|
||
/// run() 要求输出对象形状与模型完全一致（扁平 List 会被拒）。
|
||
final List<List<List<double>>> _output;
|
||
|
||
TfliteDetector._(this._interpreter, this._labels, this._numClasses,
|
||
this._numAnchors, this._output);
|
||
|
||
/// 模型缺失或加载失败返回 null（App 降级为仅预览）。
|
||
/// 在后台 isolate 内调用（模型字节由主 isolate 读取后传入）。
|
||
static Future<TfliteDetector?> fromBuffer(
|
||
Uint8List bytes, List<String> labels) async {
|
||
try {
|
||
final interpreter = Interpreter.fromBuffer(
|
||
bytes,
|
||
options: InterpreterOptions()..threads = 4,
|
||
);
|
||
return TfliteDetector._fromModel(interpreter, labels);
|
||
} catch (_) {
|
||
return null;
|
||
}
|
||
}
|
||
|
||
/// 输出布局 [1, 4+nc, anchors] 取自模型本身，类别数不与 labels 文件长度耦合。
|
||
factory TfliteDetector._fromModel(
|
||
Interpreter interpreter, List<String> labels) {
|
||
final shape = interpreter.getOutputTensor(0).shape;
|
||
final numClasses =
|
||
shape.length >= 3 && shape[1] > 4 ? shape[1] - 4 : labels.length;
|
||
final numAnchors = shape.length >= 3 && shape[2] > 0 ? shape[2] : 2100;
|
||
final output = List.generate(
|
||
1,
|
||
(_) => List.generate(
|
||
numClasses + 4,
|
||
(_) => List<double>.filled(numAnchors, 0),
|
||
),
|
||
);
|
||
return TfliteDetector._(
|
||
interpreter, labels, numClasses, numAnchors, output);
|
||
}
|
||
|
||
/// 原始数据接口（后台 isolate 用，不依赖 CameraImage）。
|
||
/// 输出坐标统一反算为原图归一化空间（与 MotionDetector 一致），
|
||
/// 否则 CENTER_CROP 裁剪偏移会让检测框系统性偏移。
|
||
List<DetectionResult> detectRaw({
|
||
required List<Uint8List> planes,
|
||
required List<int> strides,
|
||
required int width,
|
||
required int height,
|
||
required bool isBgra,
|
||
}) {
|
||
preprocess(
|
||
planes: planes,
|
||
strides: strides,
|
||
width: width,
|
||
height: height,
|
||
isBgra: isBgra);
|
||
// 传原始字节视图而非 Float32List：tflite_flutter 会对非 ByteBuffer/Uint8List
|
||
// 输入调用 resizeInputTensor（1 维 [1486848]），使 node 0 TRANSPOSE prepare 失败
|
||
_interpreter.run(_input.buffer.asUint8List(), _output);
|
||
final dets = postprocess();
|
||
// 反算与 preprocess 的 scale/dx/dy 公式一致（704 输入空间 → 原图归一化）
|
||
final scale = inputSize / width < inputSize / height
|
||
? inputSize / width
|
||
: inputSize / height;
|
||
final dx = (inputSize - width * scale) / 2;
|
||
final dy = (inputSize - height * scale) / 2;
|
||
if (dx == 0 && dy == 0) return dets;
|
||
return dets
|
||
.map((r) => r.copyWith(
|
||
left: (r.left * inputSize - dx) / (width * scale),
|
||
right: (r.right * inputSize - dx) / (width * scale),
|
||
top: (r.top * inputSize - dy) / (height * scale),
|
||
bottom: (r.bottom * inputSize - dy) / (height * scale),
|
||
))
|
||
.toList();
|
||
}
|
||
|
||
/// 按像素格式分派：iOS bgra8888 单平面 / Android yuv420 多平面。
|
||
void preprocess({
|
||
required List<Uint8List> planes,
|
||
required List<int> strides,
|
||
required int width,
|
||
required int height,
|
||
required bool isBgra,
|
||
}) {
|
||
if (isBgra) {
|
||
_preprocessBgra(planes[0], strides[0], width, height);
|
||
} else {
|
||
_preprocessYuv(planes, strides, width, height);
|
||
}
|
||
}
|
||
|
||
/// BGRA8888 单平面（iOS）：每像素 4 字节 [b,g,r,a]，双线性采样，
|
||
/// letterbox（等比缩到长边 704，短边黑边补 0，与 YOLO 训练一致）。
|
||
void _preprocessBgra(Uint8List src, int stride, int srcW, int srcH) {
|
||
final plane = inputSize * inputSize;
|
||
final scale = inputSize / srcW < inputSize / srcH
|
||
? inputSize / srcW
|
||
: inputSize / srcH;
|
||
final dx = (inputSize - srcW * scale) / 2;
|
||
final dy = (inputSize - srcH * scale) / 2;
|
||
|
||
for (var oy = 0; oy < inputSize; oy++) {
|
||
final syf = (oy - dy) / scale;
|
||
if (syf < 0 || syf >= srcH) {
|
||
for (var ox = 0; ox < inputSize; ox++) {
|
||
final p = oy * inputSize + ox;
|
||
_input[p] = 0;
|
||
_input[p + plane] = 0;
|
||
_input[p + 2 * plane] = 0;
|
||
}
|
||
continue;
|
||
}
|
||
for (var ox = 0; ox < inputSize; ox++) {
|
||
final p = oy * inputSize + ox;
|
||
final sxf = (ox - dx) / scale;
|
||
if (sxf < 0 || sxf >= srcW) {
|
||
_input[p] = 0;
|
||
_input[p + plane] = 0;
|
||
_input[p + 2 * plane] = 0;
|
||
continue;
|
||
}
|
||
final x0 = sxf.floor(), y0 = syf.floor();
|
||
final x1 = x0 < srcW - 1 ? x0 + 1 : x0;
|
||
final y1 = y0 < srcH - 1 ? y0 + 1 : y0;
|
||
final fx = sxf - x0, fy = syf - y0;
|
||
|
||
// BGRA 字节序：+0 B、+1 G、+2 R、+3 A
|
||
final i00 = y0 * stride + x0 * 4;
|
||
final i10 = y0 * stride + x1 * 4;
|
||
final i01 = y1 * stride + x0 * 4;
|
||
final i11 = y1 * stride + x1 * 4;
|
||
final r00 = src[i00 + 2].toDouble();
|
||
final g00 = src[i00 + 1].toDouble();
|
||
final b00 = src[i00].toDouble();
|
||
final r10 = src[i10 + 2].toDouble();
|
||
final g10 = src[i10 + 1].toDouble();
|
||
final b10 = src[i10].toDouble();
|
||
final r01 = src[i01 + 2].toDouble();
|
||
final g01 = src[i01 + 1].toDouble();
|
||
final b01 = src[i01].toDouble();
|
||
final r11 = src[i11 + 2].toDouble();
|
||
final g11 = src[i11 + 1].toDouble();
|
||
final b11 = src[i11].toDouble();
|
||
|
||
_input[p] = _bl(r00, r10, r01, r11, fx, fy) / 255.0;
|
||
_input[p + plane] = _bl(g00, g10, g01, g11, fx, fy) / 255.0;
|
||
_input[p + 2 * plane] = _bl(b00, b10, b01, b11, fx, fy) / 255.0;
|
||
}
|
||
}
|
||
}
|
||
|
||
/// letterbox 缩放 + YUV → RGB 归一化 0~1（NCHW），双线性采样。
|
||
/// 兼容 NV12（iOS 双平面，UV 交错）与 I420（Android 三平面）。
|
||
void _preprocessYuv(
|
||
List<Uint8List> planes, List<int> strides, int srcW, int srcH) {
|
||
final plane = inputSize * inputSize;
|
||
final y = planes[0];
|
||
final nv12 = planes.length == 2;
|
||
final uv = nv12 ? planes[1] : null;
|
||
final u = nv12 ? null : planes[1];
|
||
final v = nv12 ? null : planes[2];
|
||
final yStride = strides[0];
|
||
final uvStride = strides[1];
|
||
|
||
// U/V 平面采样（nv12：偶位 U 奇位 V；i420：三平面分离）
|
||
double uAt(int x, int y) => nv12
|
||
? uv![y * uvStride + x * 2] - 128.0
|
||
: u![y * uvStride + x] - 128.0;
|
||
double vAt(int x, int y) => nv12
|
||
? uv![y * uvStride + x * 2 + 1] - 128.0
|
||
: v![y * uvStride + x] - 128.0;
|
||
|
||
final scale = inputSize / srcW < inputSize / srcH
|
||
? inputSize / srcW
|
||
: inputSize / srcH;
|
||
final dx = (inputSize - srcW * scale) / 2;
|
||
final dy = (inputSize - srcH * scale) / 2;
|
||
|
||
for (var oy = 0; oy < inputSize; oy++) {
|
||
final syf = (oy - dy) / scale;
|
||
if (syf < 0 || syf >= srcH) {
|
||
for (var ox = 0; ox < inputSize; ox++) {
|
||
final p = oy * inputSize + ox;
|
||
_input[p] = 0;
|
||
_input[p + plane] = 0;
|
||
_input[p + 2 * plane] = 0;
|
||
}
|
||
continue;
|
||
}
|
||
for (var ox = 0; ox < inputSize; ox++) {
|
||
final p = oy * inputSize + ox;
|
||
final sxf = (ox - dx) / scale;
|
||
if (sxf < 0 || sxf >= srcW) {
|
||
_input[p] = 0;
|
||
_input[p + plane] = 0;
|
||
_input[p + 2 * plane] = 0;
|
||
continue;
|
||
}
|
||
final x0 = sxf.floor(), y0 = syf.floor();
|
||
final x1 = x0 < srcW - 1 ? x0 + 1 : x0;
|
||
final y1 = y0 < srcH - 1 ? y0 + 1 : y0;
|
||
final fx = sxf - x0, fy = syf - y0;
|
||
|
||
// Y 双线性
|
||
final y00 = y[y0 * yStride + x0].toDouble();
|
||
final y10 = y[y0 * yStride + x1].toDouble();
|
||
final y01 = y[y1 * yStride + x0].toDouble();
|
||
final y11 = y[y1 * yStride + x1].toDouble();
|
||
final yy = _bl(y00, y10, y01, y11, fx, fy);
|
||
|
||
// U/V 双线性（4:2:0 半分辨率，按像素坐标定位后除 2）
|
||
final maxUx = srcW ~/ 2 - 1;
|
||
final maxUy = srcH ~/ 2 - 1;
|
||
final ux0 = (x0 ~/ 2).clamp(0, maxUx).toInt();
|
||
final uy0 = (y0 ~/ 2).clamp(0, maxUy).toInt();
|
||
final ux1 = (x1 ~/ 2).clamp(0, maxUx).toInt();
|
||
final uy1 = (y1 ~/ 2).clamp(0, maxUy).toInt();
|
||
final u00 = uAt(ux0, uy0);
|
||
final u10 = uAt(ux1, uy0);
|
||
final u01 = uAt(ux0, uy1);
|
||
final u11 = uAt(ux1, uy1);
|
||
final uu = _bl(u00, u10, u01, u11, fx, fy);
|
||
|
||
final v00 = vAt(ux0, uy0);
|
||
final v10 = vAt(ux1, uy0);
|
||
final v01 = vAt(ux0, uy1);
|
||
final v11 = vAt(ux1, uy1);
|
||
final vv = _bl(v00, v10, v01, v11, fx, fy);
|
||
|
||
// 有限范围展开（VideoRange Y 16~235，Cb/Cr 16~240）
|
||
final yr = (yy - 16.0) * (255.0 / 219.0);
|
||
final un = uu * (255.0 / 224.0);
|
||
final vn = vv * (255.0 / 224.0);
|
||
|
||
// NCHW：r/g/b 分平面存储
|
||
_input[p] = (yr + 1.402 * vn) / 255.0;
|
||
_input[p + plane] = (yr - 0.344136 * un - 0.714136 * vn) / 255.0;
|
||
_input[p + 2 * plane] = (yr + 1.772 * un) / 255.0;
|
||
}
|
||
}
|
||
}
|
||
|
||
static double _bl(double a, double b, double c, double d, double fx,
|
||
double fy) =>
|
||
(1 - fx) * (1 - fy) * a + fx * (1 - fy) * b +
|
||
(1 - fx) * fy * c + fx * fy * d;
|
||
|
||
List<DetectionResult> postprocess() {
|
||
final out = _output[0];
|
||
final boxes = <DetectionResult>[];
|
||
for (var a = 0; a < _numAnchors; a++) {
|
||
final cx = out[0][a];
|
||
final cy = out[1][a];
|
||
final w = out[2][a];
|
||
final h = out[3][a];
|
||
var bestCls = 0;
|
||
var bestScore = 0.0;
|
||
for (var c = 0; c < _numClasses; c++) {
|
||
final s = out[4 + c][a];
|
||
if (s > bestScore) {
|
||
bestScore = s;
|
||
bestCls = c;
|
||
}
|
||
}
|
||
final label =
|
||
bestCls < _labels.length ? _labels[bestCls] : 'unknown';
|
||
// 低分池保留，供运动检测提升显示
|
||
if (bestScore < minScore) continue;
|
||
boxes.add(DetectionResult(
|
||
label: label,
|
||
score: bestScore,
|
||
left: (cx - w / 2).clamp(0.0, 1.0),
|
||
top: (cy - h / 2).clamp(0.0, 1.0),
|
||
right: (cx + w / 2).clamp(0.0, 1.0),
|
||
bottom: (cy + h / 2).clamp(0.0, 1.0),
|
||
));
|
||
}
|
||
final kept = nms(boxes, iouThreshold);
|
||
return kept.take(maxDetections).toList();
|
||
}
|
||
|
||
void dispose() => _interpreter.close();
|
||
}
|