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

Adds AO and thickness baking support #670

Open
wants to merge 19 commits into
base: main
Choose a base branch
from

Conversation

gfodor
Copy link

@gfodor gfodor commented Aug 7, 2024

This PR adds AO + Thickness map baking. The new AOThicknessMaterial replaces the AmbientOcclusionMaterial and can be set into an AO_ONLY mode to replicate the old behavior of rendering a typical ambient occlusion map. The AO_AND_THICKNESS mode (the default) bakes AO into the red channel and thickness into the green channel, as per the expected three.js aoMap and thicknessMap conventions.

A fixed up UV unwrapper based on xatlas is in UVGenerator that can be used to generate UV maps suitable for baking onto specified geometries.

A aoBake example is checked in that uses a new AOThicknessMapGenerator to bake textures using this new material. The generator first bakes the map using the typical path tracing approach, and then optionally floods in seams. The seam filling algorithm does a first pass with an island expansion algorithm (in UVEdgeExpander) and then also mip floods the map to reduce the size and deal with any residual gaps. The UVTriangleDataTextureGenerator is used to project the triangle positions and normals into UV space onto a texture for the ray tracing, this class can likely be used in general for doing UV-oriented path tracing for other things like lightmapping.

Copy link
Owner

@gkjohnson gkjohnson left a comment

Choose a reason for hiding this comment

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

Amazing! Thank you. This turned out to be quite a big PR - I think it will be easiest to go through a review component by component so I'll start with the comments on the UVUnwrapper class.

A couple more comments and questions on the example and other classes - whats the difference between "UVEdgeExpander" and "MipFlooder"? I expected only MipFlooder to be needed.

Regarding the example it's a bit hard to tell what's going on with the custom SSS shader and the current model. The aviator helmet model isn't water tight so generating a thickness map for it won't necessarily make a lot of sense or be coherent. Here are some changes that I think would be good:

  • We can change the model later (after this is merged) to something properly water tight.
  • Remove the SSS shader and just use three.js' built-in MeshPhysicalMaterial with transmission to demonstrate the thickness map
  • Add an option to toggle between transmissive material and opaque material to show off AO and thickness map
  • Add a toggle to render the generating ao map in the corner so it's easier to understand what's happening (for me as well!)
  • Lower the number of samples run per frame so the framerate isn't reduces to 3 or 4 fps during generation


generate( geometries, onProgress = null ) {

if ( geometries.constructor !== Array ) {
Copy link
Owner

Choose a reason for hiding this comment

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

We should use Array.isArray here, I think

src/utils/UVGenerator.js Show resolved Hide resolved

}

const params = { padding: 2 };
Copy link
Owner

Choose a reason for hiding this comment

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

We should expose this as an option on the class so it can be set via generator.padding = 10. It looks like this is in pixels but what does that mean if there's not texture resolution specified here?

Copy link
Author

Choose a reason for hiding this comment

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

That's a good question. I added the padding here to ensure that the edge filler has enough room to have one pixel from each adjacent island. Here are the comments in xatlas regarding the configuration of resolution, padding, and texelsPerUnit:

	// Number of pixels to pad charts with.
	uint32_t padding = 0;

	// Unit to texel scale. e.g. a 1x1 quad with texelsPerUnit of 32 will take up approximately 32x32 texels in the atlas.
	// If 0, an estimated value will be calculated to approximately match the given resolution.
	// If resolution is also 0, the estimated value will approximately match a 1024x1024 atlas.
	float texelsPerUnit = 0.0f;

	// If 0, generate a single atlas with texelsPerUnit determining the final resolution.
	// If not 0, and texelsPerUnit is not 0, generate one or more atlases with that exact resolution.
	// If not 0, and texelsPerUnit is 0, texelsPerUnit is estimated to approximately match the resolution.
	uint32_t resolution = 0

I explicitly did not set the resolution because I don't want xatlas splitting the atlas into subatlases if it decides it has run out of room: it would add a bunch of complexity to the code I didn't want to deal with. That said, when both resolution and texelsPerUnit are set to zero, I don't actually know how it could be dealing with padding in terms of pixels, since it can't know how big a pixel is. Maybe it assumes it's a 1024x1024 resolution atlas as per the estimator it uses in computing texelsPerUnit.

Copy link
Owner

Choose a reason for hiding this comment

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

From the comments it sounds like no splitting should happen as long at texelsPerUnit is set to 0. Good to know for the future, though.

Maybe it assumes it's a 1024x1024 resolution atlas as per the estimator it uses in computing texelsPerUnit.

It sounds like 1024x1024 is used. To that end it might make the most sense to parameterize this as a percentage of the UV space in a public option so it's agnostic to resolution. Ie a padding of 2 would be set as 2 / 1024 which can be converted to pixels before passing it into x atlas.

src/utils/UVGenerator.js Show resolved Hide resolved
Comment on lines +112 to +113
const newPositionArray = new Float32Array( meshData.newVertexCount * 3 );
const newNormalArray = new Float32Array( meshData.newVertexCount * 3 );
Copy link
Owner

Choose a reason for hiding this comment

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

Can you explain why these position and normal attributes are modified by xatlas? I would have expected only a new uv buffer to be generated.

Copy link
Author

Choose a reason for hiding this comment

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

From XAtlas README:

The xatlas::Atlas instance created in the first step now contains the result: each input mesh added by xatlas::AddMesh has a corresponding new mesh with a UV channel. New meshes have more vertices (the UV channel adds seams), but the same number of indices.

Comment on lines +115 to +116
let newUvArray = null, newUv2Array = null;
const useSecondChannel = this.channel === 2;
Copy link
Owner

Choose a reason for hiding this comment

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

Three.js supports uv channels from 0 to 3 (mapped to uv, uv1, uv2, uv3) so it would be good to support that here.

if ( this._module ) return;

const wasmurl = new URL( '../../node_modules/@gfodor/xatlas-web/dist/xatlas-web.wasm', import.meta.url );
this._module = XAtlas( {
Copy link
Owner

Choose a reason for hiding this comment

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

This class needs a "dispose" function to destroy the wasm handles, right?

src/utils/UVGenerator.js Outdated Show resolved Hide resolved
const geometry = geometries[ i ];

const originalVertexCount = geometry.attributes.position.count;
const originalIndexCount = geometry.index.count;
Copy link
Owner

Choose a reason for hiding this comment

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

We'll need to handle the case where geometry doesn't actually have an index.


if ( geometry.attributes.uv ) {

xatlas.HEAPF32.set( geometry.attributes.uv.array, meshInfo.uvOffset / Float32Array.BYTES_PER_ELEMENT );
Copy link
Owner

Choose a reason for hiding this comment

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

Why do we need to pass the UVs and normals into the library? I would have expected only position to be needed.

Copy link
Author

Choose a reason for hiding this comment

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

I'm not sure if it's strictly necessary - XAtlas has a mode you can run it in where you hand it UVs for each mesh and it will generate a new composed atlas using those UVs. My guess is that these aren't accessed by XAtlas for the case we are using it for.

@gfodor
Copy link
Author

gfodor commented Aug 9, 2024

Great, thanks for the review! Re: MipFlooder and UVEdgeExpander, I started out with just the MipFlooding algorithm you suggested. The problem is, that for UV edge seams it's not really a great fit. The UV maps end up having a zero alpha channel at the edges, and the mip flooding algorithm results in undesirable mis-coloring at the edges of UV islands. I took an example UV map and ran it through the python code example to confirm it wasn't my implementation, but maybe I'm missing something. A more naive edge expansion fragment shader got the job done by just expanding the edges themselves by copying neighboring pixels. The main benefit of the MipFlooding according to that blog is that it reduces the texture sizes, so I kept it in as a second pass but I don't think it's necessary. I didn't test it without it, so maybe the MipFlooding is helping with any residual gaps that don't get covered by the edge filler, which right now only does one pass. I think we could instead drop the MipFlooder and run the edge expander N times, which would flood the whole thing properly. I don't have a strong opinion, and perhaps the mip flooder will be useful when doing lightmaps since the size of those will matter a lot more.

@gfodor
Copy link
Author

gfodor commented Aug 9, 2024

Also re: the SSS I'm not totally familiar with the concept of water-tightness and the transmission property in three.js. I don't mind changing it, but the way I've been using thickness in my own work is with this subsurface scattering shader, and as you can see the corners of the base of the model are slightly highlighted (because they are thinner, and therefore light can pass through them per the shader.) I agree the model is not a good one for demonstrating transmissiveness in either case, I built most of the code with the simple box model I shared above, and then reverted to the AO model used as a full integration test since it's more complex, multi-geometry, etc.

@gfodor
Copy link
Author

gfodor commented Aug 9, 2024

(The SSS shader I put into the example is a slightly modified version of the three.js SSS example, which adds a second diffuse effect and also, imo, properly accesses the thickness in the green channel as per the three.js docs. The SSS shader in three.js uses the red channel, for some bizarre reason.)

@gkjohnson
Copy link
Owner

gkjohnson commented Aug 10, 2024

I started out with just the MipFlooding algorithm you suggested. The problem is, that for UV edge seams it's not really a great fit. The UV maps end up having a zero alpha channel at the edges, and the mip flooding algorithm results in undesirable mis-coloring at the edges of UV islands.

Gotcha. Since textures with alpha are their primary example I wonder if the alpha transparency masks any of the small artifacts you get from this approach.

I think we could instead drop the MipFlooder and run the edge expander N times, which would flood the whole thing properly. I don't have a strong opinion, and perhaps the mip flooder will be useful when doing lightmaps since the size of those will matter a lot more.

I think there's still value in filling in all the gaps so that we don't wind up with color and alpha bleeding when mip maps are generated. So lets keep it.

Also re: the SSS I'm not totally familiar with the concept of water-tightness and the transmission property in three.js.

Three.js doesn't require it but intuitively it seems necessary for something like generating a thickness map. Looking at the geometry again the topology may not be an issue for this model but it's not a super easy model to tell what's going on. We can change the model later, though.

Generally I just want to keep these examples easy to understand when looking at them and simply coded so they're easy to maintain so avoiding custom shaders as much as possible would be best. A looking at a model with SSS + AO makes both effects difficult to see. I'd like to be able to look at and evaluate the quality of each map individually.

The SSS shader in three.js uses the red channel, for some bizarre reason.

As far as I can recall that shader has been around for years - far before the transmissive material in three.js. It probably hasn't been updated to be consistent with materials.

@gkjohnson
Copy link
Owner

Hey @gfodor I have a need for automatic UV unwrapping so I may be able to put some time into working on exposing a few more options in the UV generator portion and enabling it to handle more generalized sets of geometry attributes. I didn't want to step on anything you're working on, though.

Also I'm not sure if you saw my gfodor#1 PR that updates the example primarily to make the results more clear. I'm also particularly interested in this issue relating to the UV generator stalling on 0% with that rabbit model. Perhaps it's just taking a long time because there are no UVs present? Not sure if you have more insight.

@gfodor
Copy link
Author

gfodor commented Aug 30, 2024

I went ahead and merged your changes. Two things I guess:

  • For the bunny I would try adjusting the padding, if you dig back into this before I get to it just set the padding to zero to see if it converges. (Just a guess)

  • The thing I was hitting when trying to integrate this work with a real project was all the wasm path stuff needed to be tweaked (and I landed that here), I think the current code now is probably less than ideal. How do you want wasm loading to work with three-gpu-pathtracer for the xatlas wasm binary?

@gkjohnson
Copy link
Owner

all the wasm path stuff needed to be tweaked (and I landed that here)

Are you using a build process / bundler? Which one? Unfortunately it seems these things rely heavily on the bundler (or lack of) context. A lot of bundlers now are only able to pick up on finding and including a needed static file if it's referenced via a relative URL path (same with workers). The example wasn't working with the string path that was provided before. I think we should just choose the common denominator for these features and then users can provide their own path if needed.

How do you want wasm loading to work with three-gpu-pathtracer for the xatlas wasm binary?

Referencing "node_modules" directly isn't ideal, either, in my opinion but we can deal with that another time. Unfortunately the npm, bundling, and module import ecosystem have never really caught up on these things.

@gkjohnson
Copy link
Owner

This xatlas wrapper seems very fully-featured and will soon add support for multiple atlases, as well. I may be worth switching to this one instead of maintaining a UV unwrapper in this project:

https://github.com/repalash/xatlas-three

I can look into it a bit more when I have a bit of time.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants