Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Feature] Sped Up Height Map Renderer #699

Merged
merged 3 commits into from
Feb 28, 2025
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -8,18 +8,19 @@
import com.badlogic.gdx.graphics.g3d.Renderable;
import com.badlogic.gdx.graphics.g3d.RenderableProvider;
import com.badlogic.gdx.graphics.g3d.shaders.DefaultShader;
import com.badlogic.gdx.graphics.glutils.ShaderProgram;
import com.badlogic.gdx.math.Vector2;
import com.badlogic.gdx.utils.Array;
import com.badlogic.gdx.utils.Pool;
import net.mgsx.gltf.scene3d.attributes.PBRColorAttribute;
import org.bytedeco.javacpp.BytePointer;
import org.bytedeco.javacpp.FloatPointer;
import org.bytedeco.opencv.global.opencv_core;
import org.bytedeco.opencv.opencv_core.Mat;
import org.lwjgl.opengl.GL41;
import us.ihmc.sensorProcessing.heightMap.HeightMapTools;
import us.ihmc.euclid.transform.RigidBodyTransform;
import us.ihmc.log.LogTools;
import us.ihmc.rdx.shader.RDXShader;
import us.ihmc.rdx.shader.RDXUniform;

import java.nio.FloatBuffer;

/**
* Renders a height map as a point cloud. The height map is stored as a 16-bit grayscale image.
Expand All @@ -33,31 +34,16 @@ public class RDXHeightMapRenderer implements RenderableProvider
private Renderable renderable;

public static final int FLOATS_PER_CELL = 8;
private final VertexAttributes vertexAttributes = new VertexAttributes(new VertexAttribute(VertexAttributes.Usage.Position,
3,
ShaderProgram.POSITION_ATTRIBUTE),
new VertexAttribute(VertexAttributes.Usage.ColorUnpacked,
4,
ShaderProgram.COLOR_ATTRIBUTE),
new VertexAttribute(VertexAttributes.Usage.Generic,
1,
GL41.GL_FLOAT,
false,
"a_size"));
private final RDXUniform screenWidthUniform = RDXUniform.createGlobalUniform("u_screenWidth",
(shader, inputID, renderable, combinedAttributes) -> shader.set(inputID,
shader.camera.viewportWidth));
private final RDXUniform multiColorUniform = RDXUniform.createGlobalUniform("u_multiColor", (shader, inputID, renderable, combinedAttributes) ->
{
int multiColor = 0;
shader.set(inputID, multiColor);
});

private float[] intermediateVertexBuffer;
private final VertexAttributes vertexAttributes = new VertexAttributes(new VertexAttribute(VertexAttributes.Usage.Generic, 1, "a_height"));

private int totalCells;
// Uniforms
private int centerIndex;
private final Vector2 gridCenter = new Vector2();
private float cellSize;
private float heightScalingFactor;
private float heightOffset;

public void create(int numberOfCells)
public void create(int maxCells)
{
GL41.glEnable(GL41.GL_VERTEX_PROGRAM_POINT_SIZE);

Expand All @@ -66,78 +52,103 @@ public void create(int numberOfCells)
renderable.meshPart.offset = 0;
renderable.material = new Material(PBRColorAttribute.createBaseColorFactor(Color.WHITE));

totalCells = numberOfCells;
if (renderable.meshPart.mesh != null)
renderable.meshPart.mesh.dispose();
boolean isStatic = false;
int maxIndices = 0;
renderable.meshPart.mesh = new Mesh(isStatic, totalCells, maxIndices, vertexAttributes);
renderable.meshPart.mesh = new Mesh(isStatic, maxCells, maxIndices, vertexAttributes);

RDXShader shader = new RDXShader(getClass());
shader.create();
shader.getBaseShader().register(DefaultShader.Inputs.viewTrans, DefaultShader.Setters.viewTrans);
shader.getBaseShader().register(DefaultShader.Inputs.projTrans, DefaultShader.Setters.projTrans);
shader.registerUniform(screenWidthUniform);
shader.registerUniform(multiColorUniform);
registerUniforms(shader);
shader.init(renderable);
renderable.shader = shader.getBaseShader();

LogTools.info("Vertex Buffer Size: {}", totalCells * FLOATS_PER_CELL);
intermediateVertexBuffer = new float[totalCells * FLOATS_PER_CELL];
LogTools.info("Vertex Buffer Size: {}", maxCells * FLOATS_PER_CELL);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you find this LogTools useful? I'd say we can remove it.

}

@SuppressWarnings("CodeBlock2Expr")
private void registerUniforms(RDXShader rdxShader)
{
rdxShader.getBaseShader().register(DefaultShader.Inputs.viewTrans, DefaultShader.Setters.viewTrans);
rdxShader.getBaseShader().register(DefaultShader.Inputs.projTrans, DefaultShader.Setters.projTrans);

RDXUniform screenWidthUniform = RDXUniform.createGlobalUniform("u_screenWidth", (shader, inputID, renderable, combinedAttributes) ->
{
shader.set(inputID, shader.camera.viewportWidth);
});
rdxShader.registerUniform(screenWidthUniform);

RDXUniform centerIndexUniform = RDXUniform.createGlobalUniform("u_centerIndex", (shader, inputID, renderable, combinedAttributes) ->
{
shader.set(inputID, centerIndex);
});
rdxShader.registerUniform(centerIndexUniform);

RDXUniform gridCenterUniform = RDXUniform.createGlobalUniform("u_gridCenter", (shader, inputID, renderable, combinedAttributes) ->
{
shader.set(inputID, gridCenter);
});
rdxShader.registerUniform(gridCenterUniform);

RDXUniform cellSizeUniform = RDXUniform.createGlobalUniform("u_cellSize", (shader, inputID, renderable, combinedAttributes) ->
{
shader.set(inputID, cellSize);
});
rdxShader.registerUniform(cellSizeUniform);

RDXUniform heightScalingFactorUniform = RDXUniform.createGlobalUniform("u_heightScalingFactor", (shader, inputID, renderable, combinedAttributes) ->
{
shader.set(inputID, heightScalingFactor);
});
rdxShader.registerUniform(heightScalingFactorUniform);

RDXUniform heightOffsetUniform = RDXUniform.createGlobalUniform("u_heightOffset", (shader, inputID, renderable, combinedAttributes) ->
{
shader.set(inputID, heightOffset);
});
rdxShader.registerUniform(heightOffsetUniform);
}

public void update(RigidBodyTransform zUpFrameToWorld,
BytePointer heightMapPointer,
public void update(Mat heightMapImage,
float heightOffset,
float gridCenterX,
float gridCenterY,
int centerIndex,
float cellSizeXYInMeters,
float heightScalingFactor)
{
zUpFrameToWorld.getTranslation().setZ(0);
// Update uniforms
this.heightOffset = heightOffset;
this.gridCenter.set(gridCenterX, gridCenterY);
this.centerIndex = centerIndex;
this.cellSize = cellSizeXYInMeters;
this.heightScalingFactor = heightScalingFactor;

int cellsPerAxis = 2 * centerIndex + 1;
FloatBuffer verticesBuffer = renderable.meshPart.mesh.getVerticesBuffer(true);

for (int xIndex = 0; xIndex < cellsPerAxis; xIndex++)
// Ensure correct length and initialize buffer
int cellsPerAxis = 2 * centerIndex + 1;
int totalCells = cellsPerAxis * cellsPerAxis;
if (renderable.meshPart.size != totalCells)
{
for (int yIndex = 0; yIndex < cellsPerAxis; yIndex++)
{
double xPosition = HeightMapTools.indexToCoordinate(xIndex, gridCenterX, cellSizeXYInMeters, centerIndex);
double yPosition = HeightMapTools.indexToCoordinate(yIndex, gridCenterY, cellSizeXYInMeters, centerIndex);

/* look at the header docs for decoding the height map values as below */
int heightIndex = xIndex * cellsPerAxis + yIndex;
int vertexIndex = heightIndex * FLOATS_PER_CELL;
int height = heightMapPointer.getShort(heightIndex * 2L) & 0xFFFF;
float zPosition = ((float) height / heightScalingFactor);
zPosition -= heightOffset;

intermediateVertexBuffer[vertexIndex] = (float) xPosition;
intermediateVertexBuffer[vertexIndex + 1] = (float) yPosition;
intermediateVertexBuffer[vertexIndex + 2] = zPosition;

// Color (0.0 to 1.0)
double[] redGreenBlue = HeightMapTools.getRedGreenBlue(zPosition);
Color color = new Color((float) redGreenBlue[0], (float) redGreenBlue[1], (float) redGreenBlue[2], 1.0f);
intermediateVertexBuffer[vertexIndex + 3] = color.r;
intermediateVertexBuffer[vertexIndex + 4] = color.g;
intermediateVertexBuffer[vertexIndex + 5] = color.b;
intermediateVertexBuffer[vertexIndex + 6] = color.a;

// Size
intermediateVertexBuffer[vertexIndex + 7] = 0.02f;
}
renderable.meshPart.size = totalCells;
verticesBuffer.limit(totalCells);
}

renderable.meshPart.size = totalCells;
renderable.meshPart.mesh.setVertices(intermediateVertexBuffer, 0, totalCells * FLOATS_PER_CELL);
FloatPointer verticesPointer = new FloatPointer(verticesBuffer);
Mat verticesMat = new Mat(cellsPerAxis, cellsPerAxis, opencv_core.CV_32FC1, verticesPointer);
heightMapImage.convertTo(verticesMat, opencv_core.CV_32FC1);

verticesMat.close();
verticesPointer.close();
}

@Override
public void getRenderables(Array<Renderable> renderables, Pool<Renderable> pool)
{
renderables.add(renderable);
if (renderable != null)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we create the renderable as a field = new Renderable();, we can remove this null check. And we can remove the call in the create(int maxCells) method that created the new object for the renderable. Saves a line, and removes the confusion around the null check.

renderables.add(renderable);
}

public void dispose()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -195,8 +195,7 @@ public void update()
{
if (heightMapImage.ptr(0) != null)
{
heightMapRenderer.update(zUpToWorldTransform,
heightMapImage.ptr(0),
heightMapRenderer.update(heightMapImage,
(float) RapidHeightMapManager.getHeightMapParameters().getHeightOffset(),
zUpToWorldTransform.getTranslation().getX32(),
zUpToWorldTransform.getTranslation().getY32(),
Expand Down
Original file line number Diff line number Diff line change
@@ -1,30 +1,103 @@
#type vertex
#version 410

layout(location = 0) in vec3 a_position;
layout(location = 1) in vec4 a_color;
layout(location = 2) in float a_size;
layout(location = 0) in float a_height;

out vec4 v_color;

uniform mat4 u_viewTrans;
uniform mat4 u_projTrans;
uniform float u_screenWidth;
uniform int u_multiColor;

uniform int u_centerIndex;
uniform vec2 u_gridCenter;
uniform float u_cellSize;
uniform float u_heightScalingFactor;
uniform float u_heightOffset;

float indexToCoordinate(int index, float gridCenter)
{
return (index - u_centerIndex) * u_cellSize + gridCenter;
}

float linearInterpolate(float a, float b, float alpha)
{
return (1.0f - alpha) * a + alpha * b;
}

vec4 getColor(float height)
{
// Using interpolation between key color points
float r = 0, g = 0, b = 0;
float magentaR = 1.0, magentaG = 0.0, magentaB = 1.0;
float orangeR = 1.0, orangeG = 200.0 / 255.0, orangeB = 0.0;
float yellowR = 1.0, yellowG = 1.0, yellowB = 0.0;
float blueR = 0.0, blueG = 0.0, blueB = 1.0;
float greenR = 0.0, greenG = 1.0, greenB = 0.0;
float gradientSize = 0.2;
float gradientLength = 1.0;
float alpha = mod(height, gradientLength);
if (alpha < 0)
alpha = 1 + alpha;
while (alpha > 5 * gradientSize)
alpha -= 5 * gradientSize;

if (alpha <= gradientSize * 1)
{
r = linearInterpolate(magentaR, blueR, (alpha) / gradientSize);
g = linearInterpolate(magentaG, blueG, (alpha) / gradientSize);
b = linearInterpolate(magentaB, blueB, (alpha) / gradientSize);
}
else if (alpha <= gradientSize * 2)
{
r = linearInterpolate(blueR, greenR, (alpha - gradientSize * 1) / gradientSize);
g = linearInterpolate(blueG, greenG, (alpha - gradientSize * 1) / gradientSize);
b = linearInterpolate(blueB, greenB, (alpha - gradientSize * 1) / gradientSize);
}
else if (alpha <= gradientSize * 3)
{
r = linearInterpolate(greenR, yellowR, (alpha - gradientSize * 2) / gradientSize);
g = linearInterpolate(greenG, yellowG, (alpha - gradientSize * 2) / gradientSize);
b = linearInterpolate(greenB, yellowB, (alpha - gradientSize * 2) / gradientSize);
}
else if (alpha <= gradientSize * 4)
{
r = linearInterpolate(yellowR, orangeR, (alpha - gradientSize * 3) / gradientSize);
g = linearInterpolate(yellowG, orangeG, (alpha - gradientSize * 3) / gradientSize);
b = linearInterpolate(yellowB, orangeB, (alpha - gradientSize * 3) / gradientSize);
}
else if (alpha <= gradientSize * 5)
{
r = linearInterpolate(orangeR, magentaR, (alpha - gradientSize * 4) / gradientSize);
g = linearInterpolate(orangeG, magentaG, (alpha - gradientSize * 4) / gradientSize);
b = linearInterpolate(orangeB, magentaB, (alpha - gradientSize * 4) / gradientSize);
}

return vec4(r, g, b, 1.0f);
}

void main()
{
vec4 pointInCameraFrame = u_viewTrans * vec4(a_position.x, a_position.y, a_position.z, 1);
int cellsPerAxis = 2 * u_centerIndex + 1; // width and height of the height map

int xIndex = gl_VertexID / cellsPerAxis;
int yIndex = gl_VertexID % cellsPerAxis;

float xPosition = indexToCoordinate(xIndex, u_gridCenter.x);
float yPosition = indexToCoordinate(yIndex, u_gridCenter.y);
float zPosition = (a_height / u_heightScalingFactor) - u_heightOffset;

vec4 pointInCameraFrame = u_viewTrans * vec4(xPosition, yPosition, zPosition, 1);
vec4 projectedSpriteCornerZero = u_projTrans * vec4(0.0, 0.0, pointInCameraFrame.z, pointInCameraFrame.w);

vec4 projectedSpriteCorner = u_projTrans * vec4(a_size, a_size, pointInCameraFrame.z, pointInCameraFrame.w);
vec4 projectedSpriteCorner = u_projTrans * vec4(0.02f, 0.02f, pointInCameraFrame.z, pointInCameraFrame.w);
float projectedSize = u_screenWidth * projectedSpriteCorner.x / projectedSpriteCorner.w;

gl_PointSize = 0.5 * (projectedSize);

gl_Position = u_projTrans * pointInCameraFrame;

v_color = a_color;
v_color = getColor(zPosition);
}

#type fragment
Expand Down
Loading