By using this site, you agree to our Privacy Policy and our Terms of Use. Close

Forums - PC - I trained a Deep Learning Real Time Image Quality Cleaner

Hi all, I trained a deep learning model that cleans image quality output in games, while inferencing in real time. It works after the upscaler, rather than a substitute of it. Here are some examples from Cyberpunk 2077 on an RTX 4090m. Make sure to right click -> view on a larger screen. The model is a CNN-ViT hybrid model. It has a ViT encoder and CNN decoder. While the images have 30 fps in the top left, that was a locked framerate so that the training data is synchronized. It runs at 100fps (versus 130fps w/o) on an RTX 4090m @1080p. In this example, the original is FSR 2 Performance @1080p,while the cleaned image is improved. It is a moderate improvement, I'd say it is like going up a level from performance to balanced, balanced to quality, etc., for a given upscale mode. Currently I only trained on two games (Cyberpunk 2077 and Shadow of the Tomb Raider) but I plan to train on more. The improvement is mostly spatial, it doesn't get rid of temporal shimmering and other similar artifacts, but I hope to solve those too! I have an internal app that works like Lossless Scaling (captures the input using WGC and outputs with the cleaned image.) It works only on the final output, there is no buffer data other than the color buffer. The encoder model works as a stand-in for the various buffers. There are some temporal features, but again no motion vectors. 



Around the Network

SOTR



Looks really good!



Edit: Switch to 4k in YouTube to combat compression. 

Last edited by sc94597 - 22 hours ago

That's pretty cool, although it sucks that there's any need for something like this. I'm not interested enough to learn why games keep doing this, so it's just bizarre why this blurry look is a thing in the first place.

Is this more of a hobby project, or do you perhaps have some grander plans for this tool in the future?



Around the Network
Zkuq said:

Is this more of a hobby project, or do you perhaps have some grander plans for this tool in the future?

It started as an experiment because I wanted to see if I could use feature embeddings that capture mathematical/statistical surprise in an image to detect regions of artifacts in the image, and then use those artifact-detections to help guide a relatively lightweight CNN decoder to clean up the image, mostly by masking regions of non-correction before applying corrections. My hypothesis was that areas of high surprise would be flagged as "artifacts"/"anomalous" and the decoder - given high quality examples (in this case the same image with DLAA applied rather than say FSR Performance), would adjust the "artifact" regions to more closely match the higher quality output through masking, slight color value corrections, contrast corrections, etc. The artifact-detector features would help prevent global corrections that add additional artifacts. 

The models were small, so they could train on any GPU with at least 12GB of VRAM within half a day before the networks got saturated. 

Then it pretty much evolved from there. I needed a way to test it in-game after the model was trained, since it was hard to tell subjective image quality from stills and short clips alone, so I had to basically recreate how Lossless Scaling does it using the WGC API. This was the hardest part of the work for me, because I don't have much knowledge on Windows systems programming. But eventually I was able to get it to work with CUDA and TensorRT on a cloned game Window. I also tried DirectML to support non-Nvidia GPUs but it isn't performant enough to make it worth it, at least without more effort. Will have to look at better ways to accelerate this on AMD/Intel GPUs in Windows. Maybe will try Vulkan. Also the capture doesn't work right with HDR on, so I have to figure that out. 

Currently experimenting with optical flow models to see if they'll help with temporal artifact clean-up, because I am pretty happy with the spatial clean-up, but temporal artifacts look identical to the lower-quality upscale, if not slightly worse. 

It might end up being something I open-source. If I did that, I'd definitely want to figure out the AMD/Intel hardware acceleration problem because people with those GPUs are probably who this would benefit the most. Might even look into seeing if I can get it working on NPUs that are being paired with most new systems. 



Be careful.
I have been informed by very reliable sources that using AI risks summoning a demon that will eat your soul.