Implement the PSNR vision metric

[softmaximalist]

Pull Request Template

Checklist

• Confirmed that cargo run-checks command has been executed.
• Made sure the book is up to date with changes in this PR.

Related Issues/PRs

Image quality metrics #4312

Changes

Implemented the PSNR vision metric in the crates/burn-train/src/metric/vision/psnr.rs file. It computes the per-image PSNR values and then average them across all the images in the batch. Since the PSNR (in dB) is defined as following:

$$ PSNR = 10 \cdot \log_{10}\left(\frac{MAX_{I}^2}{MSE}\right) $$

the MSE value is clamped a minimum value of epsilon to avoid division by 0. The epsilon field of the PsnrMetricConfig struct is set to 1e-10 by default. Users can set a custom/different epsilon value as well.

Testing

Added 18 test cases to the psnr.rs file. Tests cover different input shapes, different PSNR values, different image representation (where maximum possible pixel values are different), multichannel tensors, running average, setting custom name and/or epsilon, panics, etc. The test cases verify actual expected values rather than just ranges (which is approach taken by the other open PR).

Note

• Right after creating this PR, I noticed that there is another open PR also implementing PSNR (Added PSNR metric for image quality evaluation #4377).
• Below are the main differences (that I noticed) between my PR and the other PR:
  • My test cases verify very specific expected values while the other PR verifies if the psnr is within some range.
  • The other PR implements a default method for the metric which I don't. This is because the PSNR metric always needs have the max_pixel_val set but this value depends on the specific image format used by the user. Depending on the whether the images are normalized to [0, 1] range, or they are 8, 16, 32, etc bits, this value changes. Hence, in my code, the user who is aware of their specific image formats must always set this value.
  • My code only accepts 4d tensors with the shape [N, C, H, W] while the other PR accepts tensors with more than 1 dimension (dim >= 2)
    • I designed the code this way since I did not want the users to misuse the metric due to ambiguity.
    • I might be wrong but I believe the implementation in the other PR is flawed. While the code in the other PR accepts tensors with more than 1 dimension, their update() method always treats the first dimension (dim 0) as the batches. This can easily go wrong in many different cases such as when the input tensors have shapes such as [H, W], [C, H, W], etc which are still valid shapes in their code.
  • While the approach taken for computing MSE and also dealing with division by zero are different between the two PRs, they should not affect the actual functionality of the code.

[codecov]

Codecov Report

❌ Patch coverage is 99.66777% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 68.91%. Comparing base (3cd0671) to head (225686c).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/burn-train/src/metric/vision/psnr.rs	99.66%	1 Missing ⚠️

❌ Your project check has failed because the head coverage (68.91%) is below the target coverage (80.00%). You can increase the head coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4379      +/-   ##
==========================================
+ Coverage   68.86%   68.91%   +0.05%     
==========================================
  Files        1412     1413       +1     
  Lines      168245   168544     +299     
==========================================
+ Hits       115857   116153     +296     
- Misses      52388    52391       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.