AI 与模型

使用 ESRGAN-TF2 与 TFLite 在 Flutter 中实现任意尺寸图像超分

Kaggle 上的 ESRGAN-TF2 TFLite 模型输入固定为 50×50,记录张量读写、边缘补齐、分块推理与重叠裁剪去接缝的实现过程。

约 13 分钟读完
目录
  1. 1. 加载模型与配置硬件加速 Delegate
  2. 2. 单块张量转换与 runInference 实现
  3. 3. 边缘补齐 → 分块推理 → 拼接裁剪
  4. 4. 实际接入时的接缝处理与性能边界

在 Android 端基于 Flutter(配合 tflite_flutter 与纯 Dart 图像处理库 image 4.x)接入 Kaggle 上的 ESRGAN-TF2 模型 实现 4 倍图像超分辨率时,首先遇到了输入尺寸固定的限制:该 .tflite 权重的输入张量形状固定为 [1, 50, 50, 3](float32),输出张量固定为 [1, 200, 200, 3],只要传入图片的宽高不是 50×50,推理结果就会报错或输出异常画面。

最初查阅的资料大多基于原版 TensorFlow SavedModel 格式的 ESRGAN(支持动态输入宽高),因此先尝试在读取原图尺寸后调用 interpreter.resizeInputTensor 动态调整输入形状,但在该 TFLite 模型上实测不可行——导出阶段部分内部算子形状已被固化,无法直接送入整张大图。

在当时没有重训或重新导出动态轴权重的前提下,最终采用了端侧分块推理方案:先将原图右下边缘补齐至分块步长的整数倍,切分为若干 50×50 图块逐一调用 TFLite 推理(放大为 200×200),最后按网格坐标拼接还原并裁除填充边。

1. 加载模型与配置硬件加速 Delegate

按运行平台分别挂载 Android 的 GpuDelegateV2 / XNNPackDelegate 或 iOS 的 Metal GpuDelegate,加载模型后可以通过 getInputTensors / getOutputTensors 核对实际的张量形状与数据类型:

Future<void> loadModel() async {
  final options = InterpreterOptions()..threads = 4;

  if (Platform.isAndroid) {
    options.addDelegate(GpuDelegateV2());
  } else if (Platform.isIOS) {
    options.addDelegate(GpuDelegate());
  }

  interpreter = await Interpreter.fromAsset(modelPath, options: options);
  // ESRGAN-TF2 (esrgan-tf2_1.tflite):
  // input  shape = [1, 50, 50, 3], type = float32 (RGB 0.0..255.0)
  // output shape = [1, 200, 200, 3], type = float32 (RGB 0.0..255.0)
  inputTensor = interpreter.getInputTensors().first;
  outputTensor = interpreter.getOutputTensors().first;
}

(注:如果在部分老旧 Android 机型上初始化 GpuDelegateV2 抛出异常,可以在 try / catch 中降级回 XNNPackDelegate(),保证 CPU 多线程路径可用。)

2. 单块张量转换与 runInference 实现

ESRGAN-TF2 的 TFLite 权重直接接收未归一化到 [0, 1] 的原始浮点像素值(即 0.0 至 255.0 的 RGB 三通道 Float32),输出则可能因为卷积过冲出现略小于 0 或略大于 255 的浮点数。如果不做 clamp(0, 255) 直接转回整型通道,高对比度边缘很容易出现彩色溢点。

复用固定大小的 Float32List 缓冲区完成单块 50×50 → 200×200 推理的实现如下:

const int _inSize = 50;
const int _outSize = 200;

final Float32List _inputBuffer = Float32List(1 * _inSize * _inSize * 3);
final Float32List _outputBuffer = Float32List(1 * _outSize * _outSize * 3);

img.Image runInference(img.Image block) {
  var inIdx = 0;
  for (var y = 0; y < _inSize; y++) {
    for (var x = 0; x < _inSize; x++) {
      final pixel = block.getPixel(x, y);
      _inputBuffer[inIdx++] = pixel.r.toDouble();
      _inputBuffer[inIdx++] = pixel.g.toDouble();
      _inputBuffer[inIdx++] = pixel.b.toDouble();
    }
  }

  interpreter.run(_inputBuffer.buffer, _outputBuffer.buffer);

  final outImage = img.Image(width: _outSize, height: _outSize);
  var outIdx = 0;
  for (var y = 0; y < _outSize; y++) {
    for (var x = 0; x < _outSize; x++) {
      final r = _outputBuffer[outIdx++].clamp(0.0, 255.0).round();
      final g = _outputBuffer[outIdx++].clamp(0.0, 255.0).round();
      final b = _outputBuffer[outIdx++].clamp(0.0, 255.0).round();
      outImage.setPixelRgb(x, y, r, g, b);
    }
  }
  return outImage;
}

把 _inputBuffer 与 _outputBuffer 作为实例成员复用,可以避免在一张大图切出几百个图块时反复分配与回收多维 Dart List,显著减轻 GC 抖动。

3. 边缘补齐 → 分块推理 → 拼接裁剪

读取原图后,先检查宽高是否为 50 的整数倍,不足则补齐右下边缘,逐块调用 runInference 后拼接为 4 倍大图,并按原图真实宽高的 4 倍裁去填充区域:

img.Image? processImage(Uint8List rawBytes) {
  final originImage = img.decodeImage(rawBytes);
  if (originImage == null) return null;

  const inputSize = 50;
  const scale = 4;

  var paddedImage = originImage;
  if (originImage.width % inputSize != 0 || originImage.height % inputSize != 0) {
    paddedImage = padImage(originImage, inputSize);
  }

  final blocks = splitImage(paddedImage, inputSize);
  final processedBlocks = <img.Image>[
    for (final block in blocks) runInference(block),
  ];

  final combined = combineImages(
    processedBlocks,
    paddedImage.width * scale,
    paddedImage.height * scale,
    inputSize * scale,
  );

  if (paddedImage != originImage) {
    return img.copyCrop(
      combined,
      x: 0,
      y: 0,
      width: originImage.width * scale,
      height: originImage.height * scale,
    );
  }
  return combined;
}

补齐、切块与拼接辅助函数需要严格保持相同的遍历顺序(外层遍历 x、内层遍历 y),防止拼接时行列错位:

img.Image padImage(img.Image originImage, int tileSize) {
  final paddedWidth = ((originImage.width + tileSize - 1) ~/ tileSize) * tileSize;
  final paddedHeight = ((originImage.height + tileSize - 1) ~/ tileSize) * tileSize;

  final paddedImage = img.Image(width: paddedWidth, height: paddedHeight);
  // 尽量用边缘像素延展代替纯白填充,减少最右侧与最下侧图块的边界色偏
  img.compositeImage(
    paddedImage,
    originImage,
    dstX: 0,
    dstY: 0,
    dstW: originImage.width,
    dstH: originImage.height,
  );
  return paddedImage;
}

List<img.Image> splitImage(img.Image paddedImage, int tileSize) {
  final blocks = <img.Image>[];
  for (var x = 0; x < paddedImage.width; x += tileSize) {
    for (var y = 0; y < paddedImage.height; y += tileSize) {
      blocks.add(
        img.copyCrop(paddedImage, x: x, y: y, width: tileSize, height: tileSize),
      );
    }
  }
  return blocks;
}

img.Image combineImages(
  List<img.Image> blocks,
  int combinedWidth,
  int combinedHeight,
  int outTileSize,
) {
  final combinedImage = img.Image(width: combinedWidth, height: combinedHeight);
  final rowsPerCol = combinedHeight ~/ outTileSize;

  for (var x = 0; x < combinedWidth; x += outTileSize) {
    for (var y = 0; y < combinedHeight; y += outTileSize) {
      final index = (x ~/ outTileSize) * rowsPerCol + (y ~/ outTileSize);
      img.compositeImage(
        combinedImage,
        blocks[index],
        dstX: x,
        dstY: y,
        dstW: outTileSize,
        dstH: outTileSize,
      );
    }
  }
  return combinedImage;
}

4. 实际接入时的接缝处理与性能边界

把这套固定输入分块流程放到移动端页面中使用时,有两个需要处理的细节:

  1. 分块接缝与重叠裁剪(Overlap Tiling):无重叠的 50×50 直接切块在纯色背景或大面积平滑渐变(如天空、阴影过渡)区域,图块四边因为缺少邻域上下文,拼接后容易看出细微的十字网格线。如果画面对平滑度敏感,可以把每个 50×50 输入窗口拆成 中心有效区 42×42 + 四周各 4px 重叠上下文:步长按 42 滑动取块,模型输出 200×200 后裁掉四周各 4 × 4 = 16px 的边缘,只取中心 168×168 贴回目标画布。代价是总图块数会增加约 (50 / 42)² ≈ 1.42 倍,属于典型的画质与耗时折中。
  2. 后台 Isolate 与内存峰值控制:一张普通的 1000×750 图片即使不加重叠,也会切出 20 × 15 = 300 个图块,放大 4 倍后的目标画布达到 4000×3000(单张 RGBA 位图在内存中约占 48 MB)。interpreter.run 与 PNG 编码都是纯同步重计算,如果放在主 Isolate 执行会阻塞 UI 渲染导致掉帧;实际接入时应在独立后台 Isolate 中完成模型初始化与分块循环,每处理完若干块通过 SendPort 回传一次进度百分比,并限制输入原图的最大边长以防低内存机型 OOM。