stim/visualization/colormap.h

#ifndef STIM_COLORMAP_H
#define STIM_COLORMAP_H
#include <string>
#include <stdlib.h>
#ifdef __CUDACC__
#include "../cuda/error.h"
#endif
//saving an image to a file uses the CImg library
	//this currently throws a lot of "unreachable" warnings (as of GCC 4.8.2, nvcc 6.5.12)
#include "../image/image.h"
#define BREWER_CTRL_PTS 11
static float  BREWERCP[BREWER_CTRL_PTS*4] = {0.192157f, 0.211765f, 0.584314f, 1.0f,
                                      0.270588f, 0.458824f, 0.705882f, 1.0f,
                                      0.454902f, 0.678431f, 0.819608f, 1.0f,
                                      0.670588f, 0.85098f, 0.913725f, 1.0f,
                                      0.878431f, 0.952941f, 0.972549f, 1.0f,
                                      1.0f, 1.0f, 0.74902f, 1.0f,
                                      0.996078f, 0.878431f, 0.564706f, 1.0f,
                                      0.992157f, 0.682353f, 0.380392f, 1.0f,
                                      0.956863f, 0.427451f, 0.262745f, 1.0f,
                                      0.843137f, 0.188235f, 0.152941f, 1.0f,
                                      0.647059f, 0.0f, 0.14902f, 1.0f};
#ifdef __CUDACC__
texture<float4, cudaTextureType1D> cudaTexBrewer;
static cudaArray* gpuBrewer;
#endif
namespace stim{
enum colormapType {cmBrewer, cmGrayscale, cmRainbow};
static void buffer2image(unsigned char* buffer, std::string filename, unsigned int width, unsigned int height)
{
	/*unsigned char* non_interleaved = (unsigned char*)malloc(x_size * y_size * 3);
	unsigned int S = x_size * y_size;
	for(unsigned int i = 0; i < S; i++){
		non_interleaved[i + 0 * S] = buffer[i * 3 + 0];
		non_interleaved[i + 1 * S] = buffer[i * 3 + 1];
		non_interleaved[i + 2 * S] = buffer[i * 3 + 2];
	}*/
	//create an image object
	//cimg_library::CImg<unsigned char> image(non_interleaved, x_size, y_size, 1, 3);
	//image.save(filename.c_str());
    image<unsigned char> I;
    I.set_interleaved(buffer, width, height, 3);
	I.save(filename);
}
#ifdef __CUDACC__
static void initBrewer()
{
	//initialize the Brewer colormap
	//allocate CPU space
	float4 cpuColorMap[BREWER_CTRL_PTS];
	//define control rtsPoints
	cpuColorMap[0] = make_float4(0.192157f, 0.211765f, 0.584314f, 1.0f);
	cpuColorMap[1] = make_float4(0.270588f, 0.458824f, 0.705882f, 1.0f);
	cpuColorMap[2] = make_float4(0.454902f, 0.678431f, 0.819608f, 1.0f);
	cpuColorMap[3] = make_float4(0.670588f, 0.85098f, 0.913725f, 1.0f);
	cpuColorMap[4] = make_float4(0.878431f, 0.952941f, 0.972549f, 1.0f);
	cpuColorMap[5] = make_float4(1.0f, 1.0f, 0.74902f, 1.0f);
	cpuColorMap[6] = make_float4(0.996078f, 0.878431f, 0.564706f, 1.0f);
	cpuColorMap[7] = make_float4(0.992157f, 0.682353f, 0.380392f, 1.0f);
	cpuColorMap[8] = make_float4(0.956863f, 0.427451f, 0.262745f, 1.0f);
	cpuColorMap[9] = make_float4(0.843137f, 0.188235f, 0.152941f, 1.0f);
	cpuColorMap[10] = make_float4(0.647059f, 0.0f, 0.14902f, 1.0f);
	int width = BREWER_CTRL_PTS;
	int height = 0;
	// allocate array and copy colormap data
	cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc(32, 32, 32, 32, cudaChannelFormatKindFloat);
	HANDLE_ERROR(cudaMallocArray(&gpuBrewer, &channelDesc, width, height));
	HANDLE_ERROR(cudaMemcpyToArray(gpuBrewer, 0, 0, cpuColorMap, sizeof(float4)*width, cudaMemcpyHostToDevice));
	// set texture parameters
    cudaTexBrewer.addressMode[0] = cudaAddressModeClamp;
	//texBrewer.addressMode[1] = cudaAddressModeClamp;
    cudaTexBrewer.filterMode = cudaFilterModeLinear;
    cudaTexBrewer.normalized = true;  // access with normalized texture coordinates
	// Bind the array to the texture
    HANDLE_ERROR(cudaBindTextureToArray( cudaTexBrewer, gpuBrewer, channelDesc));
}
static void destroyBrewer()
{
    HANDLE_ERROR(cudaFreeArray(gpuBrewer));
}
template<class T>
__global__ static void applyBrewer(T* gpuSource, unsigned char* gpuDest, unsigned int N, T minVal = 0, T maxVal = 1)
{
	int i = blockIdx.y * gridDim.x * blockDim.x + blockIdx.x * blockDim.x + threadIdx.x;
    if(i >= N) return;
	//compute the normalized value on [minVal maxVal]
	float a = (gpuSource[i] - minVal) / (maxVal - minVal);
    //compensate for the additional space at the edges
    a *= (T)(BREWER_CTRL_PTS - 1)/(T)(BREWER_CTRL_PTS);
	//lookup the color
	float shift = (T)1/(2*BREWER_CTRL_PTS);
	float4 color = tex1D(cudaTexBrewer, a+shift);
	//float4 color = tex1D(cudaTexBrewer, a);
	gpuDest[i * 3 + 0] = 255 * color.x;
	gpuDest[i * 3 + 1] = 255 * color.y;
	gpuDest[i * 3 + 2] = 255 * color.z;
}
template<class T>
__global__ static void applyGrayscale(T* gpuSource, unsigned char* gpuDest, unsigned int N, T minVal = 0, T maxVal = 1)
{
    int i = blockIdx.y * gridDim.x * blockDim.x + blockIdx.x * blockDim.x + threadIdx.x;
    if(i >= N) return;
	//compute the normalized value on [minVal maxVal]
	float a = (gpuSource[i] - minVal) / (maxVal - minVal);
	//threshold
	if(a > 1)
        a = 1;
    if(a < 0)
        a = 0;
	gpuDest[i * 3 + 0] = 255 * a;
	gpuDest[i * 3 + 1] = 255 * a;
	gpuDest[i * 3 + 2] = 255 * a;
}
template<class T>
static void gpu2gpu(T* gpuSource, unsigned char* gpuDest, unsigned int nVals, T minVal = 0, T maxVal = 1, colormapType cm = cmGrayscale, int blockDim = 128)
{
	//This function converts a scalar field on the GPU to a color image on the GPU
	int gridX = (nVals + blockDim - 1)/blockDim;
	int gridY = 1;
    if(gridX > 65535)
    {
        gridY = (gridX + 65535 - 1) / 65535;
        gridX = 65535;
    }
    dim3 dimGrid(gridX, gridY);
	//int gridDim = (nVals + blockDim - 1)/blockDim;
	if(cm == cmGrayscale)
		applyGrayscale<<<dimGrid, blockDim>>>(gpuSource, gpuDest, nVals, minVal, maxVal);
	else if(cm == cmBrewer)
	{
		initBrewer();
		applyBrewer<<<dimGrid, blockDim>>>(gpuSource, gpuDest, nVals, minVal, maxVal);
		//HANDLE_ERROR(cudaMemset(gpuDest, 0, sizeof(unsigned char) * nVals * 3));
		destroyBrewer();
	}
}
template<class T>
static void gpu2cpu(T* gpuSource, unsigned char* cpuDest, unsigned int nVals, T minVal, T maxVal, colormapType cm = cmGrayscale)
{
    //this function converts a scalar field on the GPU to a color image on the CPU
    //first create the color image on the GPU
    //allocate GPU memory for the color image
    unsigned char* gpuDest;
    HANDLE_ERROR(cudaMalloc( (void**)&gpuDest, sizeof(unsigned char) * nVals * 3 ));
	//HANDLE_ERROR(cudaMemset(gpuSource, 0, sizeof(T) * nVals));
    //create the image on the gpu
    gpu2gpu(gpuSource, gpuDest, nVals, minVal, maxVal, cm);
	//HANDLE_ERROR(cudaMemset(gpuDest, 0, sizeof(unsigned char) * nVals * 3));
    //copy the image from the GPU to the CPU
    HANDLE_ERROR(cudaMemcpy(cpuDest, gpuDest, sizeof(unsigned char) * nVals * 3, cudaMemcpyDeviceToHost));
	HANDLE_ERROR(cudaFree( gpuDest ));
}
template<typename T>
static void gpu2image(T* gpuSource, std::string fileDest, unsigned int x_size, unsigned int y_size, T valMin, T valMax, colormapType cm = cmGrayscale)
{
	//allocate a color buffer
	unsigned char* cpuBuffer = NULL;
	cpuBuffer = (unsigned char*) malloc(sizeof(unsigned char) * 3 * x_size * y_size);
	//do the mapping
	gpu2cpu<T>(gpuSource, cpuBuffer, x_size * y_size, valMin, valMax, cm);
	//copy the buffer to an image
	buffer2image(cpuBuffer, fileDest, x_size, y_size);
	free(cpuBuffer);
}
#endif
template<class T>
static void cpuApplyBrewer(T* cpuSource, unsigned char* cpuDest, unsigned int N, T minVal = 0, T maxVal = 1)
{
    for(int i=0; i<N; i++)
    {
        //compute the normalized value on [minVal maxVal]
        float a = (cpuSource[i] - minVal) / (maxVal - minVal);
        if(a < 0) a = 0;
        if(a > 1) a = 1;
        float c = a * (float)(BREWER_CTRL_PTS-1);
        int ptLow = (int)c;
        float m = c - (float)ptLow;
        //std::cout<<m<<std::endl;
        float r, g, b;
        if(ptLow == BREWER_CTRL_PTS - 1)
        {
            r = BREWERCP[ptLow * 4 + 0];
            g = BREWERCP[ptLow * 4 + 1];
            b = BREWERCP[ptLow * 4 + 2];
        }
        else
        {
            r = BREWERCP[ptLow * 4 + 0] * (1-m) + BREWERCP[ (ptLow+1) * 4 + 0] * m;
            g = BREWERCP[ptLow * 4 + 1] * (1-m) + BREWERCP[ (ptLow+1) * 4 + 1] * m;
            b = BREWERCP[ptLow * 4 + 2] * (1-m) + BREWERCP[ (ptLow+1) * 4 + 2] * m;
        }
        cpuDest[i * 3 + 0] = 255 * r;
        cpuDest[i * 3 + 1] = 255 * g;
        cpuDest[i * 3 + 2] = 255 * b;
    }
}
template<class T>
static void cpu2cpu(T* cpuSource, unsigned char* cpuDest, unsigned int nVals, T valMin, T valMax, colormapType cm = cmGrayscale)
{
    if(cm == cmBrewer)
        cpuApplyBrewer(cpuSource, cpuDest, nVals, valMin, valMax);
    else if(cm == cmGrayscale)
    {
        int i;
        float a;
        float range = valMax - valMin;
        for(i = 0; i<nVals; i++)
        {
            //normalize to the range [valMin valMax]
            a = (cpuSource[i] - valMin) / range;
            if(a < 0) a = 0;
            if(a > 1) a = 1;
            cpuDest[i * 3 + 0] = 255 * a;
            cpuDest[i * 3 + 1] = 255 * a;
            cpuDest[i * 3 + 2] = 255 * a;
        }
    }
}
template<class T>
static void cpu2cpu(T* cpuSource, unsigned char* cpuDest, unsigned int nVals, colormapType cm = cmGrayscale, bool positive = false)
{
    //computes the max and min range automatically
    //find the largest magnitude value
    T maxVal = fabs(cpuSource[0]);
    for(int i=0; i<nVals; i++)
	{
        if(fabs(cpuSource[i]) > maxVal)
            maxVal = fabs(cpuSource[i]);
	}
    if(positive)
        cpu2cpu(cpuSource, cpuDest, nVals, (T)0, maxVal, cm);
    else
        cpu2cpu(cpuSource, cpuDest, nVals, -maxVal, maxVal, cm);
}
template<typename T>
static void cpu2image(T* cpuSource, std::string fileDest, unsigned int x_size, unsigned int y_size, T valMin, T valMax, colormapType cm = cmGrayscale)
{
    //allocate a color buffer
	unsigned char* cpuBuffer = (unsigned char*) malloc(sizeof(unsigned char) * 3 * x_size * y_size);
	//do the mapping
	cpu2cpu<T>(cpuSource, cpuBuffer, x_size * y_size, valMin, valMax, cm);
	//copy the buffer to an image
	buffer2image(cpuBuffer, fileDest, x_size, y_size);
	free(cpuBuffer);
}
template<typename T>
static void cpu2image(T* cpuSource, std::string fileDest, unsigned int x_size, unsigned int y_size, colormapType cm = cmGrayscale, bool positive = false)
{
    //allocate a color buffer
	unsigned char* cpuBuffer = (unsigned char*) malloc(sizeof(unsigned char) * 3 * x_size * y_size);
	//do the mapping
	cpu2cpu<T>(cpuSource, cpuBuffer, x_size * y_size, cm, positive);
	//copy the buffer to an image
	buffer2image(cpuBuffer, fileDest, x_size, y_size);
	free(cpuBuffer);
}
}	//end namespace colormap and rts
#endif