基于SNPE的Pix2Pix去雾APP推理输出图像错位问题问询

我正在开发一款基于Pix2Pix（类UNet）模型的安卓去雾APP，采用一加7手机摄像头采集图像，经去雾处理后在主界面展示结果，深度学习推理采用高通SNPE框架。目前遇到模型输出转Bitmap后图像错位的问题，初始错位效果如下：

推理核心代码 #

final List<String> result = new LinkedList<>();

final FloatTensor tensor = mNeuralNetwork.createFloatTensor(
        mNeuralNetwork.getInputTensorsShapes().get(mInputLayer));

Log.e("[MODEL]", "create tensor");
Bitmap smImage = Bitmap.createScaledBitmap(mImage, 1080, 720, true);

final int[] dimensions = tensor.getShape();
final boolean isGrayScale = (dimensions[dimensions.length -1] == 1);
float[] rgbBitmapAsFloat;
if (!isGrayScale) {
    rgbBitmapAsFloat = loadRgbBitmapAsFloat(smImage);
} else {
    rgbBitmapAsFloat = loadGrayScaleBitmapAsFloat(smImage);
}
tensor.write(rgbBitmapAsFloat, 0, rgbBitmapAsFloat.length);

Log.e("[MODEL]", "create tensor done!");
final Map<String, FloatTensor> inputs = new HashMap<>();
inputs.put(mInputLayer, tensor);
Log.e("[MODEL]", "create input tensor done!");

final long javaExecuteStart = SystemClock.elapsedRealtime();
final Map<String, FloatTensor> outputs = mNeuralNetwork.execute(inputs);
Log.e("[MODEL]", "model execute!");
final long javaExecuteEnd = SystemClock.elapsedRealtime();
mJavaExecuteTime = javaExecuteEnd - javaExecuteStart;
FloatTensor outputTensor = new FloatTensor() {
    @Override
    public void write(float[] floats, int i, int i1, int... ints) {

    }

    @Override
    public void write(float v, int... ints) {

    }

    @Override
    public int read(float[] floats, int i, int i1, int... ints) {
        return 0;
    }

    @Override
    public float read(int... ints) {
        return 0;
    }

    @Override
    public void release() {

    }
};
for (Map.Entry<String, FloatTensor> output : outputs.entrySet()) {
    Log.e("[MODEL]", "output_layer: " + output.getKey());
    if (output.getKey().equals(mOutputLayer)) {
        outputTensor = output.getValue();
        Log.e("[MODEL]", "output_layer: " + output.getKey() + ", shape: " +
                String.valueOf(outputTensor.getShape()[0]) + " " +
                String.valueOf(outputTensor.getShape()[1]) + " " +
                String.valueOf(outputTensor.getShape()[2]) + " " +
                String.valueOf(outputTensor.getShape()[3]) + " " );
        }
    }
return outputTensor;

SNPE FloatTensor转Bitmap代码 #

final float[] pixelsBatched = new float[tensor.getSize()];
tensor.read(pixelsBatched, 0, tensor.getSize());
Log.i("[IMAGE]", "size: " + String.valueOf(tensor.getSize()));
int w = 1080;
int h = 720;
Bitmap img = Bitmap.createBitmap(w, h, Bitmap.Config.ARGB_8888);
for (int y = 0; y < h; y++) {
    for (int x = 0; x < w; x++) {
        float r = pixelsBatched[y * w * 3 + x * 3 + 0] * 255;
        float g = pixelsBatched[y * w * 3 + x * 3 + 1] * 255;
        float b = pixelsBatched[y * w * 3 + x * 3 + 2] * 255;
        int color = ((int)r << 16) | ((int)g << 8) | (int)b | 0xFF000000;
        img.setPixel(x, y, color);
    }
}

return img;

初步排查 #

直接返回输入张量转Bitmap显示正常，说明输入处理无问题，问题出在推理环节。模型用PyTorch训练，导出ONNX测试正常，经onnx-sim简化后转SNPE的DLC模型，ONNX网络结构如下：

后续排查更新 #

更新1：调整通道读取逻辑 #

修改通道读取代码为：

int channelSize = w * h;
float r = pixelsBatched[y * w + x] * 255;
float g = pixelsBatched[y * w + x + channelSize] * 255;
float b = pixelsBatched[y * w + x + 2 * channelSize] * 255;

修改后图像效果：

更新2：SNPE工具查看结果 #

• snpe-dlv-viewer结果：
  
• 输入层信息：
  
• 输出层信息：
  

更新3：RGB转Float数组代码 #

float[] loadRgbBitmapAsFloat(Bitmap image) {
    final int[] pixels = new int[image.getWidth() * image.getHeight()];
    image.getPixels(pixels, 0, image.getWidth(), 0, 0,
            image.getWidth(), image.getHeight());

    final float[] pixelsBatched = new float[pixels.length * 3];
    for (int y = 0; y < image.getHeight(); y++) {
        for (int x = 0; x < image.getWidth(); x++) {
            final int idx = y * image.getWidth() + x;
            final int batchIdx = idx * 3;

            final float[] rgb = extractColorChannels(pixels[idx]);
            pixelsBatched[batchIdx]     = rgb[0];
            pixelsBatched[batchIdx + 1] = rgb[1];
            pixelsBatched[batchIdx + 2] = rgb[2];
        }
    }
    return pixelsBatched;
}

问题 #

请问该图像错位问题的可能原因是什么？

内容的提问来源于stack exchange，提问作者梁生珺