Online SSIM and MS-SSIM Computation

[Queuecumber]

What does this PR do?

Working on #1224
Fixes #1213

This PR fixes SSIM and MS-SSIM by allowing them to compute using running stats. This will make them actually usable in real situations.

Also I fixed several bugs that I found in these modules and changed the defaults to match pytorch-msssim

All tests should pass but to be honest I'm not sure how good those tests were anyway

I verified that this matches the known-good pytorch-mssim library to ~6 decimal places with default settings which was not the case before and I will be using this code in production in the coming week to see how it does.

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[Queuecumber]

I'm not sure what the best way to fix the mypy failure is, would appreciate some guidance on that. Explicitly annotating the variable as a Tensor didn't do it

[Queuecumber]

@SkafteNicki I could use some eyes on this, I've been using it in production and it seems to work well

[justusschock]

It seems strange to me, that compute is called in update (this kind of breaks our paradigm), although the implementation seems to be correct.

[justusschock]

Could you also add tests against pytorch-mssim as this seems to also support batching (which is lacking in sklearn)?

[Queuecumber]

Yeah I will add tests against pytorch_msssim, do you want to leave the sklearn tests in place?

BTW we found an interesting issue with pytorch_msssim while testing this PR internally. It really doesn't play nice with fp16 AMP. Torchmetrics handles it much better. Irrelevant for the test cases but thought I would point it out

[Queuecumber]

It seems strange to me, that compute is called in update (this kind of breaks our paradigm), although the implementation seems to be correct.

There should be a way to refactor it so that it looks more like the PSNR metric which does an _update call which is computing sum of squared error and then compute is more or less just reducing it.

I'm fine implementing that refactor if you want

[Queuecumber]

Just following up, my plan is to do the following

1. Rename _ssim_update to _ssim_check_inputs (open to suggestions on that name)
2. Rename _ssim_compute to _ssim_update and remove the reduction logic (but unsure of where to put the sum)
3. Add _ssim_compute which does the reduction (a little weird since in this case it's literally just calling reduce)
4. Update structural_similarity_index_measure to call these new functions correctly

[justusschock]

do you want to leave the sklearn tests in place?

Ideally yes, as this is the most trusted lib for non-batched stuff. Maybe we can then reduce the tests compared to sklearn and add the more thorough testing with the other lib, but some sklearn tests should stay.

Just following up, my plan is to do the following

Rename _ssim_update to _ssim_check_inputs (open to suggestions on that name)
Rename _ssim_compute to _ssim_update and remove the reduction logic (but unsure of where to put the sum)
Add _ssim_compute which does the reduction (a little weird since in this case it's literally just calling reduce)
Update structural_similarity_index_measure to call these new functions correctly

Sounds good to me!

[codecov]

Codecov Report

Merging #1231 (6748d7c) into master (3bf1491) will decrease coverage by 49%.
The diff coverage is 79%.

Additional details and impacted files

@@           Coverage Diff            @@
##           master   #1231     +/-   ##
========================================
- Coverage      86%     37%    -49%     
========================================
  Files         193     193             
  Lines       11372   11406     +34     
========================================
- Hits         9754    4245   -5509     
- Misses       1618    7161   +5543     

[Queuecumber]

OK I did the refactoring we discussed and added test cases for SSIM against pytorch-msssim (the MS-SSIM tests were already using it)

[SkafteNicki]

LGTM

[justusschock]

Thanks for doing. this refactor! LGTM!