修改我的 Tensorflow Lite 输出以按顺序处理和显示视频帧

Question

我有这段简单的代码（来自here），它将 TensorfFlow Lite 模型应用于输入图像（大小为 480x270），并在处理后显示结果图像。尽管生成的图像看起来与原始图像非常相似，但效果很好。现在我需要将其扩展到视频。基本上，它应该读取视频帧，而不是单个图像，像对图像所做的那样进行处理，并按顺序一个接一个地显示结果帧（而不是通过预先处理整个视频然后播放它）。我该如何修改它才能实现这一目标？

package com.example.mobedsr;

import android.content.res.AssetFileDescriptor;
import android.content.res.AssetManager;
import android.graphics.Bitmap;

import org.tensorflow.lite.DataType;
import org.tensorflow.lite.Interpreter;
import org.tensorflow.lite.gpu.CompatibilityList;
import org.tensorflow.lite.gpu.GpuDelegate;
import org.tensorflow.lite.support.common.ops.NormalizeOp;
import org.tensorflow.lite.support.image.ImageProcessor;
import org.tensorflow.lite.support.image.TensorImage;
import org.tensorflow.lite.support.image.ops.ResizeOp;
import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;

import java.io.FileInputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;


/** @brief  Super Resolution Model class
 *  @date   23/01/27
 */
public class SRModel {
    private boolean useGpu;

    public Interpreter interpreter;
    private Interpreter.Options options;
    private GpuDelegate gpuDelegate;
    private AssetManager assetManager;

    private final String MODEL_NAME = "evsrnet_x4.tflite"; //I want to change to esrgan.tflite

    SRModel(AssetManager assetManager, boolean useGpu) throws IOException {
        interpreter = null;
        gpuDelegate = null;

        this.assetManager = assetManager;
        this.useGpu = useGpu;

        // Initialize the TF Lite interpreter
        init();
    }

    private void init() throws IOException {
        options = new Interpreter.Options();

        // Set gpu delegate
        if (useGpu) {
            CompatibilityList compatList = new CompatibilityList();
            GpuDelegate.Options delegateOptions = compatList.getBestOptionsForThisDevice();
            gpuDelegate = new GpuDelegate(delegateOptions);
            options.addDelegate(gpuDelegate);
        }

        // Set TF Lite interpreter
        interpreter = new Interpreter(loadModelFile(), options);
    }

    /** @brief  Load .tflite model file to ByteBuffer
     *  @date   23/01/25
     */
    private ByteBuffer loadModelFile() throws IOException {
        AssetFileDescriptor assetFileDescriptor = assetManager.openFd(MODEL_NAME);
        FileInputStream fileInputStream = new FileInputStream(assetFileDescriptor.getFileDescriptor());

        FileChannel fileChannel = fileInputStream.getChannel();
        long startOffset = assetFileDescriptor.getStartOffset();
        long declaredLength = assetFileDescriptor.getDeclaredLength();

        return fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength);
    }

    public void run(Object a, Object b) {
        interpreter.run(a, b);
    }


    /** @brief  Prepare the input tensor from low resolution image
     *  @date   23/01/25
     */
    public TensorImage prepareInputTensor(Bitmap bitmap_lr) {
        TensorImage inputImage = TensorImage.fromBitmap(bitmap_lr);
        int height = bitmap_lr.getHeight();
        int width = bitmap_lr.getWidth();

        ImageProcessor imageProcessor = new ImageProcessor.Builder()
                .add(new ResizeOp(height, width, ResizeOp.ResizeMethod.NEAREST_NEIGHBOR))
                .add(new NormalizeOp(0.0f, 255.0f))
                .build();
        inputImage = imageProcessor.process(inputImage);

        return inputImage;
    }


    /** @brief  Prepare the output tensor for super resolution
     *  @date   23/01/25
     */
    public TensorImage prepareOutputTensor() {
        TensorImage srImage = new TensorImage(DataType.FLOAT32);
//        int[] srShape = new int[]{1080, 1920, 3};
        int[] srShape = new int[]{1920, 1080, 3};
        srImage.load(TensorBuffer.createFixedSize(srShape, DataType.FLOAT32));

        return srImage;
    }


    /** @brief  Convert tensor to bitmap image
     *  @date   23/01/25
     *  @param outputTensor super resolutioned image
     */
    public Bitmap tensorToImage(TensorImage outputTensor) {
        ByteBuffer srOut = outputTensor.getBuffer();
        srOut.rewind();

        int height = outputTensor.getHeight();
        int width = outputTensor.getWidth();

        Bitmap bmpImage = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888);
        int[] pixels = new int[width * height];

        for (int i = 0; i < width * height; i++) {
            int a = 0xFF;
            float r = srOut.getFloat() * 255.0f;
            float g = srOut.getFloat() * 255.0f;
            float b = srOut.getFloat() * 255.0f;

            pixels[i] = a << 24 | ((int) r << 16) | ((int) g << 8) | ((int) b);
        }

        bmpImage.setPixels(pixels, 0, width, 0, 0, width, height);

        return bmpImage;
    }


}

Answer 1

首先，确保模型对于任务来说足够快。视频帧速率通常为每秒 30 或 60 帧。换句话说，您的模型需要在不到 1/60 秒的时间内处理每一帧。

为了处理视频帧，您需要对其进行解码。您可以使用 MediaCodec API 来实现此目的。该文档有示例代码。

解码器的输出缓冲区包含解码后的像素数据。请注意，像素颜色通常以

YCbCr

颜色模型（俗称

YUV

）表示。更多详细信息请参见docs。在处理之前，您可能需要将它们转换为

RGB

。请参阅示例此处。

还有一个更高级别的 Media3 库，它在内部使用

MediaCodec

并处理频繁的用例（例如播放、调整大小、基本转换、HDR 支持）。您可能会在他们的演示中找到一些有用的东西。