gkjohnson / three-gpu-pathtracer

Path tracing renderer and utilities for three.js built on top of three-mesh-bvh.
https://gkjohnson.github.io/three-gpu-pathtracer/example/bundle/index.html
MIT License
1.38k stars 133 forks source link

Add denoiser #85

Open gkjohnson opened 2 years ago

gkjohnson commented 2 years ago

https://github.com/BrutPitt/glslSmartDeNoise

http://blog.gregzaal.com/2015/11/24/new-pixel-filter-type-blackman-harris/

DennisSmolek commented 4 months ago

Just to clarify...

I have GPU side input and output working for webGL and for WebGPU.

However, (in WebGL) threeJS specifically is having context conflicts with tensorflow.

This isn't an issue with WebGPU as the buffers aren't limited to the same canvas/state, just the same device.

My ultimate goal with denoising has always been this pathtracer, and with it being WebGL bound my #1 goal is to get ThreeJS and the denoiser to work together to pass data/textures on the GPU using webGL.

gkjohnson commented 4 months ago

However, (in WebGL) threeJS specifically is having context conflicts with tensorflow.

To be clear calling renderer.resetState isn't helping here? Is there state that isn't being reset by three.js correctly that could be fixed? Or is tensor flow expecting state to be remain unchanged asynchronously for a longer duration?

This isn't an issue with WebGPU as the buffers aren't limited to the same canvas/state, just the same device.

Interesting - this is good to know.

If it looks like it won't be possible to make any changes to three.js to address the WebGL issue then a WebGPU implementation seems like a fine substitute even if it's not ideal.

DennisSmolek commented 4 months ago

The code isn't perfect and I'm not ready to blast it out yet, but I have running demos of my denoiser and the rough first draft of the docs live.

Live WebGPU Demo Here Code: https://github.com/DennisSmolek/Denoiser

having vercel issues so I will have to push live demos manually once things get cleared up.

I hope to resolve my webGL issues soon then I'll focus on three/webGL and the pathtracer.

Please take a look at the demo/docs and let me know what you think.

I'll push out on twitter after a few cleanup passes, hopefully have things resolved.

DennisSmolek commented 4 months ago

However, (in WebGL) threeJS specifically is having context conflicts with tensorflow.

To be clear calling renderer.resetState isn't helping here? Is there state that isn't being reset by three.js correctly that could be fixed? Or is tensor flow expecting state to be remain unchanged asynchronously for a longer duration?

Honestly, I'm not sure. I put things on pause to get the repo up so I'll be readdressing. Essentially, when running in a threejs shared environment the model returns all 0's and no errors were thrown.

I tried multiple versions/setups of resetState as well as saving/restoring the state and program for tensorflow with Cody writing the save/restore state snippets. I'm working on a vanilla WebGL version and once I have that working I'll try again on the threeJS WebGL version. Now that it's public too when I run into errors it will be easier to share.

I honestly wonder if running OpenGL backend in chrome to speed up the pathtracer compile time (im on NVIDIA GPU) could also be playing a part.

If it looks like it won't be possible to make any changes to three.js to address the WebGL issue then a WebGPU implementation seems like a fine substitute even if it's not ideal. If you mean copy the output from three/webGL and move it to WebGPU this will totally work, however the CPU/GPU sync to get the data sucks. Like on average 1600ms.

It would be the same to go from three/webGL to the denoiser/gl and probably work (meaning devices that dont support webGPU could use it) but it would be slow...

The pure webGPU version is significantly more stable and easier to share buffers around. So once the pathtracer is on WebGPU I think it will be even better.

DennisSmolek commented 4 months ago

I wanted to tell everyone following in this thread first, ThreeJS/WebGL Version up and running!

Screenshot 2024-07-13 010811

I had to write an entire WebGLStateManager inside the denoiser on top of using renderer.resetState()

Plus getting raw WebGLTextures in/out of three is undocumented and has changed over the decade+ of threejs.

Checkout The Demo and lmk if it doesn't work or acts weird.

I'll be spending Saturday afternoon finishing up the docs and maybe trying to add more examples. Then I'll blast in out on SNS.

We've lost some speed to the tiling but as first drafts go I'm happy with it.

Edit: Oh and just a note, this isn't the best example on speed as the webGL version has to compile the first time it runs. The second + times will be exponentially faster but still around 100ms when using the AuxInputs (albedo, normal). I could have run a throw away pass (recommended in production) but it would slow the page load The "Denoise" time is also misleading as with GPU being asnyc we cant be 100%.

thomasaull commented 4 months ago

It's a small detail, but the demo image shifts a bit between the original and denoised version

DennisSmolek commented 4 months ago

It's a small detail, but the demo image shifts a bit between the original and denoised version

Yeah I’ll fix it. The data doesn’t really shift, it’s just the canvas flex/position is absolute while the preview is three images in a container.

thomasaull commented 4 months ago

@DennisSmolek A quick fix would be to add the following changes to .imgWrap:

gkjohnson commented 4 months ago

Very nice!

I had to write an entire WebGLStateManager inside the denoiser on top of using renderer.resetState()

I'm wondering if this is something that could be adjusted in three.js? Sounds like resetState may be missing some things? Unless the tensor flow package is expecting the state remain to unchanged, as well.

this isn't the best example on speed as the webGL version has to compile the first time it runs

I'm not sure if this is something you'd have control over in TF but the KHR_parallel_shader_compile could be able to help here. The page won't be blocked so you can compile on page load and be ready for denoising asap. Though I didn't see any compile time the performance trace so maybe this is already happening (granted I'm not on Windows).

The "Denoise" time is also misleading as with GPU being asnyc we cant be 100%.

With the EXT_disjoint_timer_query you should be able to get more precise GPU timing.

You may already know about these extensions but I thought I'd mention them.

DennisSmolek commented 4 months ago

So I'm about to post my version with the pathtracer.

Something is weird with the way it's handling output from the pathtracer.

First I think the pathtracer renders to it's renderTarget in HDR. This isn't a problem but I'd appreciate if someone could confirm it.

Second, either the blue noise/white noise thing mentioned here or something else in the way the pathtracer coalesces is causing issues. @gkjohnson You mention stratification, is that something that could be causing the issue?

Here is a 6spp image that fails the denoiser:

Screenshot 2024-07-17 215907

Here is a 4spp one that passes: K68HcTb

[The OIDN Docs]() mention:

The RT filter has certain limitations regarding the supported input images. Most notably, it cannot denoise images that were not rendered with ray tracing. Another important limitation is related to anti-aliasing filters. Most renderers use a high-quality pixel reconstruction filter instead of a trivial box filter to minimize aliasing artifacts (e.g. Gaussian, Blackman-Harris). The RT filter does support such pixel filters but only if implemented with importance sampling. Weighted pixel sampling (sometimes called splatting) introduces correlation between neighboring pixels, which causes the denoising to fail (the noise will not be filtered), thus it is not supported.

I've seen in other places people mention reflections/speculars that the denoiser rejects. So what I think is happening is the denoiser doesn't want to display black for missing spaces, so is either blending with the standard forward pass or is bluring in nearby data for the missing spots. The denoiser reads that as having image data and ignores it.

Here is a single image with my test output from the pathtracer, and the OIDN test image. When ran, it denoises the bottom and ignores the top: merged

I thought my denoiser was just broken, but I've since tested 20 different pre-OIDN images of many quality sizes and renderings and all of them work.

Max on twitter is using his own pathtracer with my denoiser and getting great results. (SUPER fast too, and he's using the base WebGL to canvas version)

Yi Shen's test images also run through mine just fine..

Is there some hidden setting/option I can adjust to change the pathtracers behavior?

gkjohnson commented 4 months ago

First I think the pathtracer renders to it's renderTarget in HDR. This isn't a problem but I'd appreciate if someone could confirm it.

The path tracer renders to a float 32 buffer and basically stores light intensity in linear color space which will often be larger than 1. Then when rendering to the canvas the linear colors are tone mapped to sRGB for display.

Second, either the blue noise/white noise thing mentioned here or something else in the way the pathtracer coalesces is causing issues. @gkjohnson You mention stratification, is that something that could be causing the issue?

It might be easier to help once I can see what other changes you've made to the pathtracer settings already or what the set up is - but if you're just using the default settings then yes it will be using stratified sampling with a blue noise offset. It's harder to tell from the tunnel image but the lego images are definitely using the stratified sampling.

You can swap to use white noise by changing the RANDOM_TYPE define value to 0 (PCG randomness).

It also looks like the images you're trying to denoise are being rendered at half-resolution (or some resolution less than 1) which may also be causing issues.

Regarding anti aliasing - I'll have to check on the filter again when I time but if the above issues don't help you can disable AA by changing this line to vec2 jitteredUv = vUv;.

It might also be worth taking a look at the setup that Yi Shen's three-gpu-pathtracer is using to see what's different.

DennisSmolek commented 4 months ago

The path tracer renders to a float 32 buffer and basically stores light intensity in linear color space which will often be larger than 1. Then when rendering to the canvas the linear colors are tone mapped to sRGB for display.

Ah ok. I'm already setup to convert from linear if needed. Will that bring it down to a normalized range or could it still be above 1? If using real light intensity I can adjust with a scalefactor to get the HDR in expected range but I haven't had to do that yet.

My pathtracer setup is very simple:

this.pathtracer = new WebGLPathTracer(this.renderer);
this.pathtracer.setScene(this.scene, this.camera);

// disable default fade and protections
this.pathtracer.renderToCanvas = false;
this.pathtracer.renderDelay = 100;
this.pathtracer.fadeDuration = 0;
this.pathtracer.minSamples = 0;
this.pathtracer.renderScale = 1;
// recent adjustment
this.pathtracer.multipleImportanceSampling = true;

Is there a way to adjust that material without using my own build of the pathtracer? I can pull and build my own if I need to but I was trying to avoid it in my demo.

One of those samples I was indeed using a lower render scale as a test, but this happens regardless.

I've looked at Y- Shen's version. He outputs the pathtracer to a canvas then copies that as image data into his denoiser. I could try this too but honestly I've not been able to replicate the results from his video when using the demo. Does it work better on your mac?

DennisSmolek commented 4 months ago

Ok tested on my mac both Yi-Shen's and my denoisers are significantly faster there and work. On windows, his pathtracer demo does not work, it just outputs black tiles. On mac it works very quickly which is nice.

I uploaded my (currently broken) pathtracer demo to Vercel. This seems to have fixed some things but let others still be broken.

you can see it here: https://denoiser-three-pathtracer.vercel.app/

The blown out areas make me think my handling of the input tensors is wrong and I need to adjust them. Either SRGB or some other setup.

But the noise is mostly gone .. 🤷‍♂️

This demo isn't totally stable or setup right so I wouldn't publish it.

gkjohnson commented 4 months ago

you can see it here: https://denoiser-three-pathtracer.vercel.app/

Is there anywhere to see the code?

The blown out areas make me think my handling of the input tensors is wrong and I need to adjust them. Either SRGB or some other setup.

But the noise is mostly gone .. 🤷‍♂️

Yeah the noise looks improved but it looks like inputs or something look like they may be off?

Ah ok. I'm already setup to convert from linear if needed. Will that bring it down to a normalized range or could it still be above 1?

It depends on how you're converting from linear and what to. If you mean you're converting to sRGB color space then yeah that would typically mean your values are in the range [0, 1]. I'd typically do this by using a full screen quad with a MeshBasicMaterial to handle the tonemapping to a new render target.

Is there a way to adjust that material without using my own build of the pathtracer? I can pull and build my own if I need to but I was trying to avoid it in my demo.

Not officially supported (we can adjust support later) but you can set the noise style with the following flag:

const pathTracer = new WebGLPathTracer( renderer );
pathTracer._pathTracer.material.defines.RANDOM_TYPE = 0;
DennisSmolek commented 4 months ago

Latest update: https://denoiser-three-pathtracer.vercel.app/

This version is 100% GPU based from three, to the pathtracer, to the denoiser, back to three for final rendering. SIGNIFICANTLY faster on my Mac which was already fast, 1000-2000ms on my PC but it's hard to tell because my PC averages 5-8,000ms in general.

Colors look weird IMO. And be sure to use "balanced" quality for the best quality check. Fast can be wonky.

The denoiser is running as it should but my demo is still struggling to figure out the right colorspace/tone settings of things. The colors are coming in brighter/whiter than they should and I think it's messing things up.

I'm also not 100% confident on my albedo & Normals generation

The code for this example is Here It's super messy and convoluted. Here is the quad/mixing shader in the rest of it.

How the renderer works.

Then the render loop:

  1. Warm up special render target for denoiser (initializes texture we later inject into) only runs first loop
  2. Run special Albedo/Normals MRT pass to generate Aux's for later
  3. Render Pathtracer (stop at maxSample)
  4. Render special blend fullscreen quad which has all textures and blends accordingly.
  5. If pathtracer samples is the set denoise target (default 6), run Denoiser
  6. Denoiser takes pathtraced texture, applies it to our standard fullscreenQuad and renders it out to another texture. (should be SRGB, Probably regular linear)
  7. Extract Raw WebGLTexture from textures, set color, albedo, and normals on denoiser
  8. Execute Denoiser
  9. Get denoiser output, Inject raw WebGL into Texture object (probably should be SRGB output, probably is Linear)
  10. Now that denoiser texture is set, next frame the blender (step 4) will correctly blend the denoised image over everything.

How the blender works:

  1. Start with base three.js render of the scene.
  2. Render pathtracer
  3. When pathtracer reaches minimum Samples use timer to start a blend between base render and pathtraced
  4. When denoise trigger number of samples reached start denoiser
  5. When denoiser results come back, blend between denoised and pathtraced
  6. Any user input or certain options changes, immediately reset the pathtracer, reset all blends and show base threejs render
  7. Repeat

All of the code is live, I'm 100% open to changes and criticisms. I know things need to be cleaned up and I'm totally open to suggestions for the main package or any example.

One thing to note is I will likely be moving the repo to Poimandres. They've agreed to take the project into the collective which will hopefully get more eyes on it to help it improve.

The NPM denoiser will probably stay the same because its so short.

bhouston commented 4 months ago

This is amazing! We are close! My bounty is still open BTW!

I think that there is a brightness issue? Compare the pathtraced with the denoised and the denoised is definitely all around brighter. I think that it should somehow preserve the overall intensities (e.g. the histogram should stay the same between the path traced and denoised) rather than brightening them.

Screenshot 2024-07-21 at 9 09 37 PM Screenshot 2024-07-21 at 9 09 42 PM
DennisSmolek commented 4 months ago

I think that there is a brightness issue? Compare the pathtraced with the denoised and the denoised is definitely all around brighter. I think that it should somehow preserve the overall intensities (e.g. the histogram should stay the same between the path traced and denoised) rather than brightening them.

Totally agree.

Part of me thinks it's something to do with colorSpace another thinks it's from my albedo calculations. If you look at the original "color" compared to the "albedo" pass (which is essentially just a basicMaterial you can see the over bright color.

The OIDN Docs say:

The albedo image is the feature image that usually provides the biggest quality improvement. It should contain the approximate color of the surfaces independent of illumination and viewing angle. For simple matte surfaces this means using the diffuse color/texture as the albedo. For other, more complex surfaces it is not always obvious what is the best way to compute the albedo, but the denoising filter is flexible to a certain extent and works well with differently computed albedos. Thus it is not necessary to compute the strict, exact albedo values but must be always between 0 and 1.

And goes on to explain more in detail.

If I set the debugging flag usePassthrough it creates a bypass version of the model (to test input/output handling) and it returns almost the exact pathtraced input, so it's for sure something happening withing the denoising flow.

But if I use other OIDN test images I get the expected results. So I'm fairly confident it's just something in how I'm putting the data in/expect the results

DennisSmolek commented 4 months ago

So I've improved my Albedo and normals pass including the ability to use worldNormals: image

I've also since learned that OIDN CAN accept SRGB as input with a flag but expects linear by default. Meaning the issue is likely not one about linear in/out though I'll keep testing.

Reading about brightness/colors makes me think the issue is something with tonemapping. Either the pathtracer output into the denoiser needs tonemapping applied or the same with the output of the denoiser.

I've also made some ground on speed improvements and adding better logging/debugging. There is a balance between speed/memory. The less batches the faster it runs but the bigger the memory draw.

There is a calculateBatchSize function that I haven't resolved the logic yet besides if (isMobile) return 1 which keeps the memory down the most but makes it the slowest.

I have also found we can force tf to use float16 textures for webGL which (may) make it faster. But will def help with mobile issues and may be why older models were crashing despite the lower memory reqs.

There are 4-5 other speed improvement ideas I have ranging from super experimental to absolutely will work but is super hard to execute.

Right now goals are:

  1. the Tonemapping/brightness
  2. Transparency
  3. Speed on Desktop
  4. Mobile stability
gkjohnson commented 4 months ago

Thanks for putting all this together @DennisSmolek!

I've taken a brief look at the code and these are some things that stick out to me:

The brightness differences may come from color space inconsistencies as you mention.

And lastly it should be figured out what the right way is to render the albedo and normal textures - specifically regarding transparent surfaces and depth of field. It's possible these textures are supposed to be averaged results in the same way the path traced input texture is. This would require changes to the path tracer, of course. From other discussions it sounds like switching to use the white noise setting in the path tracer may help the results, as well.


Whenever someone wants to make a PR into this project to add support for one of the OIDN pathtracers to WebGLPathTracer that would be great. We can start with the three.js rasterized albedo and normal textures and improve over time.

DennisSmolek commented 4 months ago

Made some changes and have some results: https://denoiser-three-pathtracer.vercel.app/ Updated On the left is what was up until a few minutes ago, on the right are with my changes.

From what I can tell, in general, the more samples the closer to ground truth.

Minor color space changes noticeable even in the pathtraced results. With the older version the bad results get magnified and blow things out. Newer version blends much closer to the original PT results.

At 100 samples, the blown out results lessen and it gets much closer to the PT source. However the newer version looks pretty consistent (although naturally better) to the 6 sample and PT version.

These are both with "fast" quality on purpose which gets the worse results but helped show the changes.

Play with it and lmk. be sure to up the samples and "denoise at"

Negatives: I did notice some roughness in dark edges. I'm not sure of the source. I think it may be the new normals. Will add a test to go between the worldspace and viewspace.

DennisSmolek commented 4 months ago
  • Albedo texture is perform 4x4 samples of the texture map which will cause a blurry texture surface.

I did this because docs said inputs should be anti-aliased the same. I didn't notice any bluryness in the albedo output (you can change the output and see) If you think it best raw we can totally remove it.

  • Looking at the Linear sRGB path traced texture conversion to sRGB - it looks as though no conversion to sRGB is performed. ~The Albedo texture, however, is rendered to sRGB~ (I'm actually not exactly sure what color space the Albedo textures are in since a RawShaderMaterial is being used) so there's a mismatch between the value ranges and (possibly) color spaces of the albedo and path tracer input. I'm not sure right this second TBH. It's the default so I thought it was coming out as Linear. I will check though.

I have made adjustments to the Albedo and Normals pass to add the worldspace normals option (to match the OIDN examples) although OIDN says either world or viewspace is acceptable. I noticed some darkness/oddness in the edges though so I think this a big area of attention.

  • The rendered normal buffer doesn't use any normal maps.

Like sampling from a provided map? I generated the normals at RT but sampling a map if it has one makes sense.

And lastly it should be figured out what the right way is to render the albedo and normal textures

Yeah, I just tried my best here and tried to keep it simple. It's straightforward to get them with threejs forward passes but (especially with transparency) may need more advanced versions. Totally open to whatever yall think best.

Sidenote, If we use the pathtracer for the Aux Inputs its highly recommended to pre-pass denoise them before sending to the final denoiser.

Whenever someone wants to make a PR into this project to add support for one of the OIDN pathtracers to WebGLPathTracer that would be great. We can start with the three.js rasterized albedo and normal textures and improve over time.

Do you have any opinions on API's/presets? Do you want to assume the pathtracer will be generating the albedos and normals by default or leave that as an external object the user has to setup/use?

Do you want the main render call to use the denoiser automatically?

I was thinking something like:

 //defaults
pathtracer.denoiseAt = 6;
pathtracer.generateAux= true;
// later
pathtracer.renderSample();
// optional
pathtracer.denoiseSample(); // denoises current sample
pathtracer.setAux(albedoTexture, normalTexture);
gkjohnson commented 4 months ago

Do you want to assume the pathtracer will be generating the albedos and normals by default or leave that as an external object the user has to setup/use?

I think the path tracer should handle all of this internally to make it as simple as possible to use. I think the API can be as simple as this:

pathtracer.maxSamples = 10 // defaults to 20 or so?
pathtracer.enableDenoiser = true // defaults to true?

The path tracer will stop rendering once maxSamples is reached regardless of whether the denoiser is enabled. And if the denoiser is enabled then denoising will occur once the cap is hit.

Setting enableDenoiser to false and maxSamples to Infinity gets the current behavior of the pathtracer. I'm torn on whether the denoiser should default to true or not. But that can be decided later.

edit

I did this because docs said inputs should be anti-aliased the same. I didn't notice any bluryness in the albedo output (you can change the output and see) If you think it best raw we can totally remove it.

Also, sampling 4x4 pixels on the texture surface (with an incompatible resolution since the sampling isn't happening in screen space) isn't the same as performing anti aliasing. You can enable hardware anti aliasing at triangle edges by setting the WebGLRenderTarget.samples field.

You can see in the current albedo output that there's no AA happening at the moment:

image