cuda标准差拉伸

标准差拉伸（tif影像波段值类型由16bit转为8bit）cuda实现版本

使用gdal2.4.4，cuda10.1 ，thrust库（计算波段均值、方差值）

1. 使用 gdal2.4.4 读取 GTiff 格式影像，读取数据至数组

2. 使用 thrust库计算 最大值、最小值、波段均值、方差等

3. cuda10.1 核函数执行条件判断赋值

头文件引用

• thrust计算最大值、最小值引用 

#include "thrust/extrema.h"

• 设备指针

#include "thrust/device_vector.h"

• thrust 可以在 cpu 和 gpu 端执行

#include "thrust/execution_policy.h"

通过调用函数的第一个参数指定 thrust::reduce(thrust::host, 或 thrust:device

• 在累加求和时注意总和值类型，数组类型为 unsigned short ，求和后会远远超过 该类型最大值，故在

auto band_sum = thrust::reduce(ptr, ptr + size, (ull)0); 指定计算类型为 unsigned long long

代码如下：

#include "cuda_runtime.h"
#include "device_launch_parameters.h"

#include "thrust/host_vector.h"
#include "thrust/device_vector.h"
#include "thrust/extrema.h"
#include "thrust/reduce.h"
#include "thrust/functional.h"
#include "thrust/execution_policy.h"

#include "gdal_util.h"
#include "cpl_conv.h"

// 求方差
struct variance : std::unary_function<us, double>
{
    variance(double m) : mean(m) { }
    const double mean;
    __host__ __device__ double operator()(us data) const
    {
        return std::pow(data - mean, 2.0);
    }
};

__global__ void pixels_std(us* data, uc* res, const ull size, us band_max, us band_min, us uc_max, us uc_min, float k, float b)
{
    ui tid = threadIdx.x + blockDim.x * blockIdx.x;
    if (tid >= size) return;

    const us d = data[tid];
    us v;
    if (d == band_min)
        v = band_min;
    else if (d <= uc_min)
        v = band_min;
    else if (d >= uc_max)
        v = band_max;
    else if (k * d + b < band_min)
        v = band_min;
    else if (k * d + b > band_max)
        v = band_max;
    else if (k * d + b > band_min && k * d + b < band_max)
        v = k * d + b;
    else
        v = d;

    res[tid] = static_cast<uc>(v);
}

int main(int argc, char* argv[])
{
    // 16bit 转 8bit
    GDALAllRegister();

    char psz_filename[1024] = "D:\\cuda\\PAN31.TIF";
    char psz_filename_new[1024] = "D:\\cuda\\PANNew.TIF";
    // GetGDALDriverManager()->AutoLoadDrivers();
    GDALDriver* tifDriver = GetGDALDriverManager()->GetDriverByName("GTiff");
    raster_info ri;

    CPLSetConfigOption("GDAL_FILENAME_IS_UTF8", "NO");
    GDALDatasetH dataset_uint16 = GDALOpen(psz_filename, GA_Update);
    
    // if (dataset_uint16 == NULL)
    get_raster_info(dataset_uint16, &ri);
    //  新影像
    GDALDataset* dataset_uint8 = tifDriver->Create(psz_filename_new, ri.width, ri.height, GDALGetRasterCount(dataset_uint16), GDT_Byte,NULL);
    dataset_uint8->SetGeoTransform(ri.geo_transform);
    dataset_uint8->SetProjection(ri.projection);

    printf("Size is %dx%dx%d\n",
        ri.width,
        ri.height,
        GDALGetRasterCount(dataset_uint8));
    printf("Pixel Size = (%.6f,%.6f)\n",
        ri.geo_transform[1], ri.geo_transform[5]);

    cudaError_t status;
    GDALRasterBandH h_band;
    GDALRasterBandH h_band2;
    const int x_size = ri.width;
    const int y_size = ri.height;
    const ull size = x_size * y_size;
    const ull malloc_size = sizeof(us) * x_size * y_size;

    // 原影像
    us* h_data;
    // 新影像
    uc* h_res;
    h_data = (us*)CPLMalloc(malloc_size);
    h_res = (uc*)CPLMalloc(size);
    us* d_data;
    uc* d_res;
    status = cudaMalloc((void**)&d_data, malloc_size);
    status = cudaMalloc((void**)&d_res, size);
    for (int i = 0; i < 3; ++i)
    {
        h_band = GDALGetRasterBand(dataset_uint16, i + 1);
        h_band2 = GDALGetRasterBand(dataset_uint8, i + 1);
        GDALRasterIO(h_band, GF_Read, 0, 0, x_size, y_size,
            h_data, x_size, y_size, GDT_UInt16, 0, 0);
        status = cudaMemcpy(d_data, h_data, malloc_size, cudaMemcpyHostToDevice);
        thrust::device_ptr<us> ptr(d_data);
        // 数组越界时抛出 msg:extrema failed to synchronize
        // 最大值最小值仅为测试
        const auto max_iter = thrust::max_element(ptr, ptr + size);
        const auto min_iter = thrust::min_element(ptr, ptr + size);
        us band_max = *max_iter;
        us band_min = *min_iter;
        band_max = 255;
        band_min = 0;
        // cpu 执行
        // auto band_sum_cpu = thrust::reduce(thrust::host, h_data, h_data + size, (ull)0);
        // gpu 执行
        // 此处总和值类型使用 unsigned long long 
        auto band_sum = thrust::reduce(ptr, ptr + size, (ull)0);
        double band_mean = band_sum / (double)size;
        // 方差 (val-mean)*(val-mean)
        auto band_std2 = thrust::transform_reduce(ptr, ptr + size, variance(band_mean), (double)0, thrust::plus<double>());
        double band_std = std::sqrt(band_std2/(double)(size-1));
        // 2.5倍标准差
        float kn = 2.5;
        float uc_max = band_mean + kn * band_std;
        float uc_min = band_mean - kn * band_std;
        float k = (band_max - band_min) / (uc_max - uc_min);
        float b = (uc_max * band_min - uc_min * band_max) / (uc_max - uc_min);
        if (uc_min <= 0)
            uc_min = 0;
            
        const ui block_size = 128;
        const ui grid_size = (size - 1) / block_size + 1;
        pixels_std << <grid_size, block_size >> > (d_data, d_res, size, band_max, band_min, uc_max, uc_min, k, b);

        cudaDeviceSynchronize();

        cudaMemcpy(h_res, d_res, size, cudaMemcpyDeviceToHost);
        //
        GDALRasterIO(h_band2, GF_Write, 0, 0, x_size, y_size,
            h_res, x_size, y_size, GDT_Byte,0, 0);
    }
    cudaFree(d_data);
    cudaFree(d_res);
    CPLFree(h_data);
    CPLFree(h_res);
    
    GDALClose(dataset_uint16);
    GDALClose(dataset_uint8);
    
    return 0;
}

python版本 使用 gdal+numpy实现，GitHub链接：

python实现版本