#include <osg/Texture2D>
#include <osg/BindImageTexture>
#include <osg/DispatchCompute>
#include <osg/Geode>
#include <osgDB/ReadFile>
#include <osgGA/StateSetManipulator>
#include <osgViewer/Viewer>
#include <osgViewer/ViewerEventHandlers>
static const char* computeSrc = {
"#version 430\n"
"precision highp float;\n"
"precision highp int;\n"
"uniform int radius;\n"
"uniform float width;\n"
"uniform float height;\n"
"const float pi = 3.1415926;\n"
"float sigma = float(radius) * 0.25;\n"
"float s = 2 * sigma * sigma;\n"
"layout (rgba32f, binding =0) highp uniform image2D uImageIn;\n"
"layout (rgba32f, binding =1) highp uniform image2D uImageOut;\n"
"layout (local_size_x = 16, local_size_y = 16, local_size_z = 1) in;\n"
"void main() {\n"
" ivec2 id = ivec2(gl_GlobalInvocationID.xy); \n"
" ivec2 size = imageSize(uImageOut); \n"
" if (id.x >= size.x || id.y >= size.y) { \n"
" return; \n"
" } \n"
" vec2 scale = vec2(1.0) / vec2(width,height); \n"
" vec4 pixel = vec4(0.0); \n"
" float weightSum = 0.0; \n"
" float weight = 0; \n"
" vec2 offset = vec2(0.0); \n"
" for (int i = -radius / 2; i < radius / 2; i++) \n"
" { \n"
" for (int j = -radius / 2; j < radius / 2; j++) \n"
" { \n"
" offset = vec2(i, j); \n"
" weight = exp(-(offset.x * offset.x + offset.y * offset.y) / s) / (pi * s); \n"
" pixel += imageLoad(uImageIn, ivec2(x, y) + scale * offset) * weight; \n"
" weightSum += weight; \n"
" } \n"
" } \n"
" pixel /= weightSum; \n"
" imageStore(uImageOut, id, pixel); \n"
" \n"
"}\n"
};
int main(int argc, char** argv)
{
osg::ArgumentParser arguments(&argc, argv);
// Create the texture as both the output of compute shader and the input of a normal quad
osg::ref_ptr<osg::Texture2D> tex2D = new osg::Texture2D;
osg::Image* pImage = osgDB::readImageFile("Images/man.jpg");
pImage->setDataVariance(osg::Object::DYNAMIC);
tex2D->setImage(pImage);
tex2D->setTextureSize(512, 512);
tex2D->setFilter(osg::Texture2D::MIN_FILTER, osg::Texture2D::LINEAR);
tex2D->setFilter(osg::Texture2D::MAG_FILTER, osg::Texture2D::LINEAR);
tex2D->setInternalFormat(GL_RGB32F_ARB);
tex2D->setSourceFormat(GL_RGB);
tex2D->setSourceType(GL_FLOAT);
// So we can use 'image2D' in the compute shader
osg::ref_ptr<osg::BindImageTexture> imagbinding = new osg::BindImageTexture(0, tex2D.get(), osg::BindImageTexture::READ_WRITE, GL_RGB32F_ARB);
osg::ref_ptr<osg::Texture2D> tex2D2 = new osg::Texture2D;
tex2D2->setTextureSize(512, 512);
tex2D2->setFilter(osg::Texture2D::MIN_FILTER, osg::Texture2D::LINEAR);
tex2D2->setFilter(osg::Texture2D::MAG_FILTER, osg::Texture2D::LINEAR);
tex2D2->setInternalFormat(GL_RGB32F_ARB);
tex2D2->setSourceFormat(GL_RGB);
tex2D2->setSourceType(GL_FLOAT);
// So we can use 'image2D' in the compute shader
osg::ref_ptr<osg::BindImageTexture> imagbinding2 = new osg::BindImageTexture(0, tex2D2.get(), osg::BindImageTexture::READ_WRITE, GL_RGB32F_ARB);
// The compute shader can't work with other kinds of shaders
// It also requires the work group numbers. Setting them to 0 will disable the compute shader
osg::ref_ptr<osg::Program> computeProg = new osg::Program;
computeProg->addShader(new osg::Shader(osg::Shader::COMPUTE, computeSrc));
// Create a node for outputting to the texture.
// It is OK to have just an empty node here, but seems inbuilt uniforms like osg_FrameTime won't work then.
// TODO: maybe we can have a custom drawable which also will implement glMemoryBarrier?
osg::ref_ptr<osg::Node> sourceNode = new osg::DispatchCompute(512 / 16, 512 / 16, 1);
sourceNode->setDataVariance(osg::Object::DYNAMIC);
sourceNode->getOrCreateStateSet()->setAttributeAndModes(computeProg.get());
sourceNode->getOrCreateStateSet()->addUniform(new osg::Uniform("uImageIn", (int)0));
sourceNode->getOrCreateStateSet()->setTextureAttributeAndModes(0, tex2D.get());
//
osg::Uniform* radius0 = new osg::Uniform("radius", 15);
sourceNode->getOrCreateStateSet()->addUniform(radius0);
osg::Uniform* width0 = new osg::Uniform("width", (float)pImage->s());
sourceNode->getOrCreateStateSet()->addUniform(width0);
osg::Uniform* height0 = new osg::Uniform("height", (float)pImage->t());
sourceNode->getOrCreateStateSet()->addUniform(height0);
//
sourceNode->getOrCreateStateSet()->addUniform(new osg::Uniform("uImageOut", (int)1));
sourceNode->getOrCreateStateSet()->setTextureAttributeAndModes(1, tex2D2.get());
// Display the texture on a quad. We will also be able to operate on the data if reading back to CPU side
osg::Geometry* geom = osg::createTexturedQuadGeometry(
osg::Vec3(), osg::Vec3(1.0f, 0.0f, 0.0f), osg::Vec3(0.0f, 0.0f, 1.0f));
osg::ref_ptr<osg::Geode> quad = new osg::Geode;
quad->addDrawable(geom);
quad->getOrCreateStateSet()->setMode(GL_LIGHTING, osg::StateAttribute::OFF);
quad->getOrCreateStateSet()->setTextureAttributeAndModes(0, tex2D2.get());
quad->getOrCreateStateSet()->setAttributeAndModes(imagbinding2.get());
// Create the scene graph and start the viewer
osg::ref_ptr<osg::Group> scene = new osg::Group;
scene->addChild(sourceNode);
scene->addChild(quad.get());
osgViewer::Viewer viewer;
viewer.addEventHandler(new osgGA::StateSetManipulator(viewer.getCamera()->getOrCreateStateSet()));
viewer.addEventHandler(new osgViewer::StatsHandler);
viewer.addEventHandler(new osgViewer::WindowSizeHandler);
viewer.setSceneData(scene.get());
return viewer.run();
}