Add TCM (CVPR 2023) + CCA (NeurIPS 2024) with Family-1 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

@@ -74,12 +74,14 @@ def __init__(
         channel_context: Mapping[str, nn.Module],
         *,
         groups: List[int],
+        max_support_slices: int = -1,
         **kwargs,
     ):
         super().__init__()
         self._kwargs = kwargs
         self.groups = list(groups)
         self.groups_acc = list(accumulate(self.groups, initial=0))
+        self.max_support_slices = int(max_support_slices)
         self.channel_context = nn.ModuleDict(channel_context)
         self.latent_codec = nn.ModuleDict(latent_codec)
 
@@ -137,8 +139,12 @@ 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)
+        # Spatial dims are the trailing two entries of any per-group shape;
+        # the channel total is determined by ``self.groups`` (so this works
+        # for both leaves that report ``(C, H, W)`` -- e.g. CheckerboardLatentCodec --
+        # and leaves that report ``(H, W)`` -- e.g. GaussianConditionalLatentCodec).
+        spatial = tuple(shape[0])[-2:]
+        y_hat = torch.zeros((n, sum(self.groups), *spatial), device=side_params.device)
         y_hat_ = y_hat.split(self.groups, dim=1)
 
         for k in range(len(self.groups)):
@@ -165,5 +171,12 @@ 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 = self._select_support(k, y_hat_)
+        ch_ctx_params = self.channel_context[f"y{k}"](self.merge_y(*support))
         return self.merge_params(ch_ctx_params, side_params)
+
+    def _select_support(self, k: int, y_hat_: List[Tensor]) -> List[Tensor]:
+        prior = list(y_hat_[:k])
+        if self.max_support_slices < 0:
+            return prior
+        return prior[: self.max_support_slices]

compressai/losses/__init__.py

@@ -28,10 +28,12 @@
 # ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 from . import pointcloud
+from .cca import CCARateDistortionLoss
 from .pointcloud import *
 from .rate_distortion import RateDistortionLoss
 
 __all__ = [
     *pointcloud.__all__,
+    "CCARateDistortionLoss",
     "RateDistortionLoss",
 ]