I'm trying to run the example on a 2160x1080 .png file but I'm getting this error:
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
Cell In[4], line 12
10 jpeg_quality: Tensor = torch.tensor([2.0])
11 # Perform differentiable JPEG coding
---> 12 image_coded: Tensor = diff_jpeg_coding(image_rgb=image, jpeg_quality=jpeg_quality)
File ~/miniconda3/lib/python3.9/site-packages/diff_jpeg/jpeg.py:49, in diff_jpeg_coding(image_rgb, jpeg_quality, quantization_table_y, quantization_table_c, ste)
47 _, _, H, W = image_rgb.shape # type: int, int, int, int
48 # Perform encoding
---> 49 y_encoded, cb_encoded, cr_encoded = jpeg_encode(
50 image_rgb=image_rgb,
51 jpeg_quality=jpeg_quality,
52 quantization_table_c=quantization_table_c,
53 quantization_table_y=quantization_table_y,
54 rounding_function=rounding_function,
55 floor_function=floor_function,
56 clipping_function=clipping_function,
57 ) # type: Tensor, Tensor, Tensor
58 image_rgb_jpeg: Tensor = jpeg_decode(
59 input_y=y_encoded,
60 input_cb=cb_encoded,
(...)
68 clipping_function=clipping_function,
69 )
70 # Clip coded image
File ~/miniconda3/lib/python3.9/site-packages/diff_jpeg/encode.py:142, in jpeg_encode(image_rgb, jpeg_quality, rounding_function, floor_function, clipping_function, quantization_table_y, quantization_table_c)
140 input_y, input_cb, input_cr = chroma_subsampling(image_ycbcr) # type: Tensor, Tensor, Tensor
141 # Patchify, DCT, and rounding
--> 142 input_y, input_cb, input_cr = patchify_8x8(input_y), patchify_8x8(input_cb), patchify_8x8(input_cr)
143 dct_y, dct_cb, dct_cr = dct_8x8(input_y), dct_8x8(input_cb), dct_8x8(input_cr) # type: Tensor, Tensor, Tensor
144 y_encoded: Tensor = quantize(
145 dct_y,
146 jpeg_quality,
(...)
150 clipping_function=clipping_function,
151 )
File ~/miniconda3/lib/python3.9/site-packages/diff_jpeg/encode.py:22, in patchify_8x8(input)
20 B, H, W = input.shape # type: int, int, int
21 # Patchify to shape [B, N, H // 8, W // 8]
---> 22 output: Tensor = input.view(B, H // 8, 8, W // 8, 8).permute(0, 1, 3, 2, 4).reshape(B, -1, 8, 8)
23 return output
RuntimeError: shape '[1, 67, 8, 135, 8]' is invalid for input of size 583200