necla-ml / diff-jpeg Goto Github PK

---------------------------------------------------------------------------
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

B, N, _, _ = input.shape  # type: int, int, int, int