Add Jpeg Compression augmentation

[gau-nernst]

🚀 The feature

Add Jpeg Compression as an augmentation.

Motivation, pitch

Jpeg compression augmentation can be useful to simulate compression artifacts from JPEG.

Since torchvision already has facilities to do Jpeg encoding and decoding from Tensor, implementation for Jpeg compression augmentation is simple.

Proof of concept

import torch
from torch import Tensor, nn
from torchvision.io import decode_jpeg, encode_jpeg


def apply_jpeg(x: Tensor, quality: int) -> Tensor:
    return decode_jpeg(encode_jpeg(x, quality))


class RandomApplyJpeg(nn.Module):
    def __init__(self, min_quality: int, max_quality: int) -> None:
        assert 1 <= min_quality <= max_quality <= 100
        super().__init__()
        self.min_quality = min_quality
        self.max_quality = max_quality

    def forward(self, x: Tensor) -> Tensor:
        quality = torch.randint(self.min_quality, self.max_quality + 1, ()).item()
        return apply_jpeg(x, quality)

Alternatives

No response

Additional context

No response

[NicolasHug]

Thanks for the feature request @gau-nernst . This sounds reasonable, let me know if you'd like to submit a PR for this.

Instead of having 2 parameters min_quality and max_quality, maybe we can have a single quality parameter that can either be an int (non-random) or a tuple (random) where quality[0] is min_quality and quality[1] is max_quality

[gau-nernst]

@NicolasHug I can try to submit a PR for this one also. Do you have any preference on how we should name it?

[NicolasHug]

Just JPEG()? Let's bikeshed on the name later :)

[NicolasHug]

@gau-nernst do you think you'll have time to submit a PR before the end of week? I'm asking because our branch cut for the next release will be on Monday, so ideally we should get this in before Monday in order for the transform to be available for torchvision 0.18.

I don't mean to rush you though, as I know you're already working on #8295. LMK and if you can't make it no worries, I'll try to do it myself.

[gau-nernst]

@NicolasHug I can do it. I made a mistake of not making a new branch for #8295 in my fork, so I wanted to get #8295 merged before working on another PR. But it's alright, I can still do it.