Unity中通过shader实现OpenCV的Remap功能

我们希望对输入图片根据Distort程序进行畸变处理：

void Distort(const float &x_src,
             const float &y_src,
             const cv::Mat &control_points,
             const float &r2,
             const cv::Mat &W,
             float &x_dst,
             float &y_dst) {
    
    int pts_num = control_points.rows;

    x_dst = 0;
    y_dst = 0;
    for (int i = 0; i < pts_num; i++) {
        float x_diff = x_src - control_points.at<cv::Vec2f>(i)[0];
        float y_diff = y_src - control_points.at<cv::Vec2f>(i)[1];

        float kernel = 1.f / sqrt(x_diff * x_diff + y_diff * y_diff + r2);

        x_dst += kernel * W.at<float>(i, 0);
        y_dst += kernel * W.at<float>(i, 1);
    }

    x_dst += (W.at<float>(pts_num, 0) + W.at<float>(pts_num + 1, 0) * x_src + W.at<float>(pts_num + 2, 0) * y_src);
    y_dst += (W.at<float>(pts_num, 1) + W.at<float>(pts_num + 1, 1) * x_src + W.at<float>(pts_num + 2, 1) * y_src);
}

输入图像	处理后图像

可以根据Distort程序，生成OpenCV中cv::remap函数所需的map1，map2映射矩阵，然后再使用cv::remap对输入图片进行处理。我们希望在Unity中，实现类似于cv::remap函数的功能。

首先我们根据Distort()生成像素点的LUT映射矩阵，并且保存到screen_calibration_lut.bin文件中：

void GenerateLut(const cv::Mat &control_points,
                 const float &r2,
                 const cv::Mat &W,
                 const int width,
                 const int height,
                 cv::Mat &lut) {
    // Generate Texcoords LUT
    // -------------------------------------------------------------------------------------------
    // By default, gl_FragCoord assumes a lower-left origin for window coordinates and assumes pixel centers are located at half-pixel coordinates. For example, the (x, y) location (0.5, 0.5) is returned for the lower-left-most pixel in a window"
    lut = cv::Mat::zeros(height, width, CV_32FC2);
    float x_distort, y_distort, x_cv, y_cv, x_tex, y_tex;
    // float x_frag, y_frag;
    for (int y = 0; y < height; y++) {
        // y_frag = y + 0.5f;
        // y_cv = float(img_height) - y_frag - 0.5f;
        y_cv = float(height) - y - 1;
        for (int x = 0; x < width; x++) {
            // x_frag = x + 0.5f;
            // x_cv = x_frag - 0.5f;
            x_cv = x;

            Distort(x_cv, y_cv, control_points, r2, W, x_distort, y_distort);
            // std::cout<<"["<<x<<", "<<y<<"] -> ["<<x_distort<<", "<<y_distort<<"]"<<std::endl;

            x_tex = (x_distort + 0.5f) / float(width);
            y_tex = 1.f - (y_distort + 0.5f) / float(height);

            lut.at<cv::Vec2f>(y, x)[0] = x_tex;
            lut.at<cv::Vec2f>(y, x)[1] = y_tex;
        }
    }
    // -------------------------------------------------------------------------------------------
}

int main(int argc, char *argv[]) {
    // ...
    cv::Mat lut;
    GenerateLut(control_points, r2, weights, screen_size.width, screen_size.height, lut); 

    std::string screen_calibration_lut_file = calibration_result_save_path + "/screen_calibration_lut.bin";

    std::ofstream out(screen_calibration_lut_file, std::ios::out | std::ios::binary | std::ios::trunc);
    if (!lut.isContinuous()) {lut = lut.clone();}
    out.write((char *)lut.data, 2 * lut.cols * lut.rows * sizeof(float));
    out.close();
    // ...
}

在Unity端中，读取screen_calibration_lut.bin文件，并且通过Unity Shader屏幕特效函数OnRenderImage()¹，在图像渲染完成后对图像进行Remap处理：

using System.Collections;
using System.Collections.Generic;
using System.IO;
using UnityEngine;
using Newtonsoft.Json;
using Newtonsoft.Json.Linq;

public class ScreenUndistortion : MonoBehaviour
{
    public bool enable_screen_undistortion = true;
    public string screen_calibration_file = "screen_calibration_lut.bin";
    public int screen_width = 1920;
    public int screen_height = 1080;

    private Material material;


    void Start()
    {
        if (enable_screen_undistortion == false) return;
        
        // ReadCalibrationData
        // ------------------------------------------------------------------------
        float[] lut = new float[screen_width * screen_height * 2];
        if (!ReadCalibrationData(screen_calibration_file, lut))
        {
            Debug.Log("Can not open file: " + screen_calibration_file);
            return;
        }
        // ------------------------------------------------------------------------

        // material
        // ------------------------------------------------------------------------
        Shader shader = Shader.Find("Hidden/ScreenUndistortion");
        if (shader == null) {
            Debug.Log("Can not found shader!");
			return;
		}

        material = new Material(shader);
    
        Texture2D tex = new Texture2D(screen_width, screen_height, TextureFormat.RGFloat, true);
        // public void SetPixelData(NativeArray<T> data, int mipLevel, int sourceDataStartIndex = 0);
        // - data	Data array to initialize texture pixels with.
        // - mipLevel	Mip level to fill.
        // - sourceDataStartIndex	Index in the source array to start copying from (default 0).
        tex.SetPixelData(lut, 0, 0);
        // public void Apply(bool updateMipmaps = true, bool makeNoLongerReadable = false);
        // - updateMipmaps	When set to true, mipmap levels are recalculated.
        // - makeNoLongerReadable	When set to true, Unity discards the copy of pixel data in CPU-addressable memory after this operation.
        tex.Apply(false, false);
        material.SetTexture("_LutTex", tex);

        MeshRenderer meshRenderer = gameObject.AddComponent<MeshRenderer>();      
        meshRenderer.material = material;

    }

    bool ReadCalibrationData(string file,
                             float[] lut)
    {
        if (lut == null || !File.Exists(file))
        {
            return false;
        }

        using(BinaryReader reader = new BinaryReader(File.Open(file, FileMode.Open)))
        {
            for(int i = 0; i < lut.Length; i++)
            {
                lut[i] = reader.ReadSingle();
            }
        }

        // Debug.Log(lut[0] + "," + lut[1] + "," + lut[2] + "," + lut[3]);

        return true;
    }

    void OnRenderImage(RenderTexture src, RenderTexture dest)
    {
        if (enable_screen_undistortion == false || material == null)
        {
            Graphics.Blit(src, dest);
        }
        else
        {
            Graphics.Blit(src, dest, material);
        }
    }

}

其中，所使用的shader文件ScreenUndistortion.shader为：

Shader "Hidden/ScreenUndistortion"
{
    Properties
    {
        _MainTex ("Texture", 2D) = "white" {}
        _LutTex("Texture", 2D) = "white" {}
    }
    SubShader
    {
        // No culling or depth
        Cull Off ZWrite Off ZTest Always

        Pass
        {
            CGPROGRAM
            #pragma vertex vert
            #pragma fragment frag

            #include "UnityCG.cginc"

            struct appdata
            {
                float4 vertex : POSITION;
                float2 uv : TEXCOORD0;
            };

            struct v2f
            {
                // Unity stores UVs in 0-1 space. [0,0] represents the bottom-left corner of the texture, and [1,1] represents the top-right. Values are not clamped; you can use values below 0 and above 1 if needed.
                float2 uv : TEXCOORD0;
                float4 vertex : SV_POSITION;
            };

            v2f vert (appdata v)
            {
                v2f o;
                o.vertex = UnityObjectToClipPos(v.vertex);
                o.uv = v.uv;
                return o;
            }

            sampler2D _MainTex;
            sampler2D _LutTex;

            fixed4 frag (v2f i) : SV_Target
            {
                float4 uv_distort = tex2D(_LutTex, i.uv);

                float4 col;

                if (uv_distort.x <= 0.0f || uv_distort.y <= 0.0f || uv_distort.x >= 1.0f || uv_distort.y >= 1.0f)
                {
                    col = float4(0, 0, 0, 1);
                }
                else
                {
                    col = tex2D(_MainTex, float2(uv_distort.x, uv_distort.y));
                } 
                
                return col;
            }
            ENDCG
        }
    }
}

---

1. https://docs.unity3d.com/ScriptReference/MonoBehaviour.OnRenderImage.html↩︎