Add TCM (CVPR 2023) + CCA (NeurIPS 2024) with codec refactor

compressai/latent_codecs/__init__.py

@@ -29,7 +29,6 @@
 
 from .base import LatentCodec
 from .channel_groups import ChannelGroupsLatentCodec
-from .channel_slice import ChannelSliceLatentCodec
 from .checkerboard import CheckerboardLatentCodec
 from .entropy_bottleneck import EntropyBottleneckLatentCodec
 from .gain import GainHyperLatentCodec, GainHyperpriorLatentCodec
@@ -41,7 +40,6 @@
 __all__ = [
     "LatentCodec",
     "ChannelGroupsLatentCodec",
-    "ChannelSliceLatentCodec",
     "CheckerboardLatentCodec",
     "EntropyBottleneckLatentCodec",
     "GainHyperLatentCodec",

compressai/latent_codecs/channel_groups.py

@@ -28,7 +28,7 @@
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 from itertools import accumulate
-from typing import Any, Dict, List, Mapping, Tuple
+from typing import Any, Dict, List, Mapping, Optional, Tuple
 
 import torch
 import torch.nn as nn
@@ -74,12 +74,18 @@ def __init__(
         channel_context: Mapping[str, nn.Module],
         *,
         groups: List[int],
+        support_slices: Optional[List[List[int]]] = None,
         **kwargs,
     ):
         super().__init__()
         self._kwargs = kwargs
         self.groups = list(groups)
         self.groups_acc = list(accumulate(self.groups, initial=0))
+        if support_slices is None:
+            support_slices = [range(k) for k in range(len(self.groups))]
+        assert len(support_slices) == len(self.groups)
+        assert all(all(0 <= j < k for j in s) for k, s in enumerate(support_slices))
+        self.support_slices = [tuple(support_slice) for support_slice in support_slices]
         self.channel_context = nn.ModuleDict(channel_context)
         self.latent_codec = nn.ModuleDict(latent_codec)
 
@@ -121,14 +127,14 @@ def compress(self, y: Tensor, side_params: Tensor) -> Dict[str, Any]:
 
         return {
             "strings": [s for ss in y_strings_groups for s in ss],
-            "shape": [y_out["shape"] for y_out in y_out_],
+            "shape": y.shape[1:],
             "y_hat": y_hat,
         }
 
     def decompress(
         self,
         strings: List[List[bytes]],
-        shape: List[Tuple[int, ...]],
+        shape: Tuple[int, ...],
         side_params: Tensor,
         **kwargs,
     ) -> Dict[str, Any]:
@@ -137,15 +143,14 @@ def decompress(
         strings_per_group = len(strings) // len(self.groups)
 
         y_out_ = [{}] * len(self.groups)
-        y_shape = (sum(s[0] for s in shape), *shape[0][1:])
-        y_hat = torch.zeros((n, *y_shape), device=side_params.device)
+        y_hat = torch.zeros((n, *shape), device=side_params.device)
         y_hat_ = y_hat.split(self.groups, dim=1)
 
         for k in range(len(self.groups)):
             params = self._get_ctx_params(k, side_params, y_hat_)
             y_out_[k] = self.latent_codec[f"y{k}"].decompress(
                 strings[strings_per_group * k : strings_per_group * (k + 1)],
-                shape[k],
+                (self.groups[k], *shape[1:]),
                 params,
             )
             y_hat_[k][:] = y_out_[k]["y_hat"]
@@ -165,5 +170,6 @@ def _get_ctx_params(
     ) -> Tensor:
         if k == 0:
             return side_params
-        ch_ctx_params = self.channel_context[f"y{k}"](self.merge_y(*y_hat_[:k]))
+        support = [y_hat_[i] for i in self.support_slices[k]]
+        ch_ctx_params = self.channel_context[f"y{k}"](self.merge_y(*support))
         return self.merge_params(ch_ctx_params, side_params)

compressai/latent_codecs/checkerboard.py

@@ -276,7 +276,7 @@ def decompress(
         assert all(len(x) == n for x in y_strings_)
 
         c, h, w = shape
-        y_i_shape = (h, w // 2)
+        y_i_shape = (c, h, w // 2)
         y_hat_ = side_params.new_zeros((2, n, c, h, w // 2))
         side_params_ = self.unembed(side_params)
 

compressai/latent_codecs/entropy_bottleneck.py

@@ -80,14 +80,16 @@ def forward(self, y: Tensor) -> Dict[str, Any]:
         return {"likelihoods": {"y": y_likelihoods}, "y_hat": y_hat}
 
     def compress(self, y: Tensor) -> Dict[str, Any]:
-        shape = y.size()[-2:]
+        shape = y.shape[1:]
         y_strings = self.entropy_bottleneck.compress(y)
-        y_hat = self.entropy_bottleneck.decompress(y_strings, shape)
+        y_hat = self.entropy_bottleneck.decompress(y_strings, shape[1:])
+        assert y_hat.shape[1:] == shape
         return {"strings": [y_strings], "shape": shape, "y_hat": y_hat}
 
     def decompress(
-        self, strings: List[List[bytes]], shape: Tuple[int, int], **kwargs
+        self, strings: List[List[bytes]], shape: Tuple[int, ...], **kwargs
     ) -> Dict[str, Any]:
         (y_strings,) = strings
-        y_hat = self.entropy_bottleneck.decompress(y_strings, shape)
+        y_hat = self.entropy_bottleneck.decompress(y_strings, shape[1:])
+        assert y_hat.shape[1:] == shape
         return {"y_hat": y_hat}