Low Rank Image Compression

Compress a batch of greyscale images with a truncated singular value decomposition (SVD), and keep only as many singular values as each image needs.

For one image A, the SVD is A = U @ diag(S) @ Vh, with the singular values in S sorted from largest to smallest. Keeping the first k of them gives the best rank k approximation A_k. Its relative error follows from the singular values alone:

error(k) = ||A - A_k||_F / ||A||_F
         = sqrt(S[k]^2 + S[k+1]^2 + ...) / sqrt(S[0]^2 + S[1]^2 + ...)

For each image, find the smallest k >= 1 with error(k) <= max_error, and return that rank together with the rank k reconstruction. Different images in the batch can need different ranks. Do the whole batch with one call to torch.linalg.svd, and use a mask on the singular values to keep a different number for each image.

Function signature

def compress(images, max_error):

Example

Input

images = torch.diag(torch.tensor([3.0, 4.0, 0.0]))[None]   # shape (1, 3, 3)
compress(images, max_error=0.7)

Output

(tensor([1]),
 tensor([[[0., 0., 0.],
          [0., 4., 0.],
          [0., 0., 0.]]]))

Reasoning

The singular values are 4, 3, 0, so the total energy is 16 + 9 = 25. Keeping one value leaves 9, and sqrt(9 / 25) = 0.6, which is below 0.7. The best rank 1 approximation keeps only the direction with value 4.

Constraints and notes

• images has shape (B, m, n), is float32 or float64, and no image is all zeros.
• Return ranks as an int64 tensor of shape (B,), and the reconstruction with the same shape and dtype as images.
• 0 <= max_error < 1. Do not change the input.