Skip to content

Commit 0288983

Browse files
YiozolmYiozolm
committed
Add MLIC series models
1 parent 81e018a commit 0288983

23 files changed

Lines changed: 5068 additions & 1 deletion

‎compressai/latent_codecs/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
from .gaussian_conditional import GaussianConditionalLatentCodec
3636
from .hyper import HyperLatentCodec
3737
from .hyperprior import HyperpriorLatentCodec
38+
from .multi_context_checkerboard import MultiContextCheckerboardLatentCodec
3839
from .rasterscan import RasterScanLatentCodec
3940

4041
__all__ = [
@@ -47,5 +48,6 @@
4748
"GaussianConditionalLatentCodec",
4849
"HyperLatentCodec",
4950
"HyperpriorLatentCodec",
51+
"MultiContextCheckerboardLatentCodec",
5052
"RasterScanLatentCodec",
5153
]

‎compressai/latent_codecs/_checkerboard_helpers.py‎

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
# Copyright (c) 2021-2025, InterDigital Communications, Inc
2+
# All rights reserved.
3+
4+
# Redistribution and use in source and binary forms, with or without
5+
# modification, are permitted (subject to the limitations in the disclaimer
6+
# below) provided that the following conditions are met:
7+
8+
# * Redistributions of source code must retain the above copyright notice,
9+
# this list of conditions and the following disclaimer.
10+
# * Redistributions in binary form must reproduce the above copyright notice,
11+
# this list of conditions and the following disclaimer in the documentation
12+
# and/or other materials provided with the distribution.
13+
# * Neither the name of InterDigital Communications, Inc nor the names of its
14+
# contributors may be used to endorse or promote products derived from this
15+
# software without specific prior written permission.
16+
17+
# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY
18+
# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
19+
# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT
20+
# NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
21+
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
22+
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
23+
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
24+
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
25+
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
26+
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
27+
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
28+
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
29+
30+
"""Pure functional helpers shared by checkerboard latent codecs.
31+
32+
These are extracted from :class:`CheckerboardLatentCodec` so that sibling
33+
codecs (e.g. :class:`MultiContextCheckerboardLatentCodec`) can reuse the
34+
exact same checkerboard split / merge / mask logic without duplicating it.
35+
A single source of truth here also means an anchor-parity boundary fix
36+
applies to every checkerboard codec at once.
37+
"""
38+
39+
from __future__ import annotations
40+
41+
import torch
42+
43+
from torch import Tensor
44+
45+
__all__ = [
46+
"embed",
47+
"embed_step",
48+
"mask_all",
49+
"mask_all_but_step",
50+
"merge",
51+
"step_parity",
52+
"unembed",
53+
"write_step",
54+
]
55+
56+
57+
def step_parity(step: str, anchor_parity: str) -> str:
58+
"""Resolve a ``step`` ('anchor' / 'non_anchor') to a parity string."""
59+
if step == "anchor":
60+
return anchor_parity
61+
if step == "non_anchor":
62+
return "odd" if anchor_parity == "even" else "even"
63+
raise ValueError(f'Invalid "step" value "{step}"')
64+
65+
66+
def unembed(y: Tensor, *, anchor_parity: str) -> Tensor:
67+
"""Separate single tensor into two even/odd checkerboard chunks.
68+
69+
.. code-block:: none
70+
71+
■ □ ■ □ ■ ■ □ □
72+
□ ■ □ ■ ---> ■ ■ □ □
73+
■ □ ■ □ ■ ■ □ □
74+
"""
75+
n, c, h, w = y.shape
76+
y_packed = y.new_zeros((2, n, c, h, w // 2))
77+
if anchor_parity == "even":
78+
y_packed[0, ..., 0::2, :] = y[..., 0::2, 0::2]
79+
y_packed[0, ..., 1::2, :] = y[..., 1::2, 1::2]
80+
y_packed[1, ..., 0::2, :] = y[..., 0::2, 1::2]
81+
y_packed[1, ..., 1::2, :] = y[..., 1::2, 0::2]
82+
else:
83+
y_packed[0, ..., 0::2, :] = y[..., 0::2, 1::2]
84+
y_packed[0, ..., 1::2, :] = y[..., 1::2, 0::2]
85+
y_packed[1, ..., 0::2, :] = y[..., 0::2, 0::2]
86+
y_packed[1, ..., 1::2, :] = y[..., 1::2, 1::2]
87+
return y_packed
88+
89+
90+
def embed(y_packed: Tensor, *, anchor_parity: str) -> Tensor:
91+
"""Combine two even/odd checkerboard chunks into single tensor.
92+
93+
.. code-block:: none
94+
95+
■ ■ □ □ ■ □ ■ □
96+
■ ■ □ □ ---> □ ■ □ ■
97+
■ ■ □ □ ■ □ ■ □
98+
"""
99+
num_chunks, n, c, h, w_half = y_packed.shape
100+
assert num_chunks == 2
101+
y = y_packed.new_zeros((n, c, h, w_half * 2))
102+
if anchor_parity == "even":
103+
y[..., 0::2, 0::2] = y_packed[0, ..., 0::2, :]
104+
y[..., 1::2, 1::2] = y_packed[0, ..., 1::2, :]
105+
y[..., 0::2, 1::2] = y_packed[1, ..., 0::2, :]
106+
y[..., 1::2, 0::2] = y_packed[1, ..., 1::2, :]
107+
else:
108+
y[..., 0::2, 1::2] = y_packed[0, ..., 0::2, :]
109+
y[..., 1::2, 0::2] = y_packed[0, ..., 1::2, :]
110+
y[..., 0::2, 0::2] = y_packed[1, ..., 0::2, :]
111+
y[..., 1::2, 1::2] = y_packed[1, ..., 1::2, :]
112+
return y
113+
114+
115+
def embed_step(
116+
step_index: int, y_i: Tensor, width: int, *, anchor_parity: str
117+
) -> Tensor:
118+
"""Embed a per-step half-width tensor back into a full-grid tensor."""
119+
n, c, h, _ = y_i.shape
120+
y_packed = y_i.new_zeros((2, n, c, h, width // 2))
121+
y_packed[step_index] = y_i
122+
return embed(y_packed, anchor_parity=anchor_parity)
123+
124+
125+
def write_step(dest: Tensor, src: Tensor, step: str, *, anchor_parity: str) -> None:
126+
"""Copy ``src`` pixels at the current step's positions into ``dest`` in-place."""
127+
parity = step_parity(step, anchor_parity)
128+
if parity == "even":
129+
dest[..., 0::2, 0::2] = src[..., 0::2, 0::2]
130+
dest[..., 1::2, 1::2] = src[..., 1::2, 1::2]
131+
else:
132+
dest[..., 0::2, 1::2] = src[..., 0::2, 1::2]
133+
dest[..., 1::2, 0::2] = src[..., 1::2, 0::2]
134+
135+
136+
def mask_all_but_step(y: Tensor, step: str, *, anchor_parity: str) -> Tensor:
137+
"""Keep only pixels in the current step, and zero out the rest."""
138+
y = y.clone()
139+
parity = step_parity(step, anchor_parity)
140+
if parity == "even":
141+
y[..., 0::2, 1::2] = 0
142+
y[..., 1::2, 0::2] = 0
143+
else:
144+
y[..., 0::2, 0::2] = 0
145+
y[..., 1::2, 1::2] = 0
146+
return y
147+
148+
149+
def mask_all(y: Tensor) -> Tensor:
150+
"""Return a zero tensor with the same shape, dtype and device as ``y``."""
151+
return torch.zeros_like(y)
152+
153+
154+
def merge(*args: Tensor) -> Tensor:
155+
"""Concatenate tensors along the channel dimension."""
156+
return torch.cat(args, dim=1)

‎compressai/latent_codecs/_selective_checkerboard.py‎

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
# Copyright (c) 2021-2025, InterDigital Communications, Inc
2+
# All rights reserved.
3+
4+
# Redistribution and use in source and binary forms, with or without
5+
# modification, are permitted (subject to the limitations in the disclaimer
6+
# below) provided that the following conditions are met:
7+
8+
# * Redistributions of source code must retain the above copyright notice,
9+
# this list of conditions and the following disclaimer.
10+
# * Redistributions in binary form must reproduce the above copyright notice,
11+
# this list of conditions and the following disclaimer in the documentation
12+
# and/or other materials provided with the distribution.
13+
# * Neither the name of InterDigital Communications, Inc nor the names of its
14+
# contributors may be used to endorse or promote products derived from this
15+
# software without specific prior written permission.
16+
17+
# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY
18+
# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
19+
# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT
20+
# NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
21+
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
22+
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
23+
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
24+
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
25+
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
26+
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
27+
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
28+
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
29+
30+
from typing import Any, Dict, List, Optional, Tuple
31+
32+
import torch
33+
import torch.nn as nn
34+
35+
from torch import Tensor
36+
37+
from compressai.entropy_models import GaussianConditional
38+
39+
from . import _checkerboard_helpers as _ckb
40+
41+
__all__ = [
42+
"apply_selective_y_hat",
43+
"apply_selective_y_hat_packed",
44+
"apply_selective_compression",
45+
"apply_selective_decompression",
46+
"compress_selected",
47+
"decompress_selected",
48+
"selective_mask",
49+
"selective_mask_packed",
50+
]
51+
52+
53+
def selective_mask(
54+
selective_predictor: Optional[nn.Module],
55+
step: str,
56+
side_params: Tensor,
57+
scales: Tensor,
58+
means: Tensor,
59+
*,
60+
anchor_parity: str,
61+
) -> Optional[Tensor]:
62+
if selective_predictor is None:
63+
return None
64+
selective_map = selective_predictor(
65+
side_params=side_params,
66+
scales=scales,
67+
means=means,
68+
step=step,
69+
)
70+
if isinstance(selective_map, dict):
71+
selective_map = selective_map["selective_map"]
72+
if selective_map.shape != scales.shape:
73+
selective_map = selective_map.expand_as(scales)
74+
if selective_map.dtype == torch.bool:
75+
mask = selective_map
76+
else:
77+
mask = selective_map >= 0.5
78+
return _ckb.mask_all_but_step(mask, step, anchor_parity=anchor_parity)
79+
80+
81+
def selective_mask_packed(
82+
selective_predictor: Optional[nn.Module],
83+
step_index: int,
84+
step: str,
85+
side_params: Tensor,
86+
scales: Tensor,
87+
means: Tensor,
88+
*,
89+
anchor_parity: str,
90+
) -> Optional[Tensor]:
91+
if selective_predictor is None:
92+
return None
93+
width = side_params.shape[-1]
94+
scales_full = _ckb.embed_step(
95+
step_index, scales, width, anchor_parity=anchor_parity
96+
)
97+
means_full = _ckb.embed_step(step_index, means, width, anchor_parity=anchor_parity)
98+
mask = selective_mask(
99+
selective_predictor,
100+
step,
101+
side_params,
102+
scales_full,
103+
means_full,
104+
anchor_parity=anchor_parity,
105+
)
106+
if mask is None:
107+
return None
108+
return _ckb.unembed(mask, anchor_parity=anchor_parity)[step_index]
109+
110+
111+
def apply_selective_y_hat(
112+
step: str,
113+
y_hat: Tensor,
114+
means: Tensor,
115+
selective_mask: Optional[Tensor],
116+
*,
117+
anchor_parity: str,
118+
) -> Tensor:
119+
if selective_mask is None:
120+
return y_hat
121+
y_hat = torch.where(selective_mask, y_hat, means)
122+
return _ckb.mask_all_but_step(y_hat, step, anchor_parity=anchor_parity)
123+
124+
125+
def apply_selective_y_hat_packed(
126+
y_hat: Tensor,
127+
means: Tensor,
128+
selective_mask: Optional[Tensor],
129+
) -> Tensor:
130+
if selective_mask is None:
131+
return y_hat
132+
return torch.where(selective_mask, y_hat, means)
133+
134+
135+
def apply_selective_compression(
136+
latent_codec: Any,
137+
y: Tensor,
138+
params: Tensor,
139+
scales: Tensor,
140+
means: Tensor,
141+
selective_mask: Optional[Tensor],
142+
) -> Dict[str, Any]:
143+
if selective_mask is None:
144+
return latent_codec.compress(y, params)
145+
return compress_selected(
146+
latent_codec.gaussian_conditional, y, scales, means, selective_mask
147+
)
148+
149+
150+
def apply_selective_decompression(
151+
latent_codec: Any,
152+
strings: List[bytes],
153+
shape: Tuple[int, ...],
154+
params: Tensor,
155+
scales: Tensor,
156+
means: Tensor,
157+
selective_mask: Optional[Tensor],
158+
) -> Dict[str, Any]:
159+
if selective_mask is None:
160+
return latent_codec.decompress([strings], shape, params)
161+
y_hat = decompress_selected(
162+
latent_codec.gaussian_conditional, strings, scales, means, selective_mask
163+
)
164+
assert y_hat.shape[1:] == shape
165+
return {"y_hat": y_hat}
166+
167+
168+
def compress_selected(
169+
gaussian_conditional: GaussianConditional,
170+
y: Tensor,
171+
scales: Tensor,
172+
means: Tensor,
173+
selective_mask: Tensor,
174+
) -> Dict[str, Any]:
175+
indexes = gaussian_conditional.build_indexes(scales)
176+
y_strings = []
177+
y_hat = means.clone()
178+
179+
for sample_index in range(y.shape[0]):
180+
mask = selective_mask[sample_index].reshape(-1)
181+
if not mask.any():
182+
y_strings.append(b"")
183+
continue
184+
185+
y_i = y[sample_index].reshape(-1)[mask].unsqueeze(0)
186+
indexes_i = indexes[sample_index].reshape(-1)[mask].unsqueeze(0)
187+
means_i = means[sample_index].reshape(-1)[mask].unsqueeze(0)
188+
y_string = gaussian_conditional.compress(y_i, indexes_i, means_i)[0]
189+
y_hat_i = gaussian_conditional.decompress([y_string], indexes_i, means=means_i)
190+
y_hat[sample_index].reshape(-1)[mask] = y_hat_i.reshape(-1).to(y_hat.dtype)
191+
y_strings.append(y_string)
192+
193+
return {"strings": [y_strings], "shape": y.shape[2:4], "y_hat": y_hat}
194+
195+
196+
def decompress_selected(
197+
gaussian_conditional: GaussianConditional,
198+
strings: List[bytes],
199+
scales: Tensor,
200+
means: Tensor,
201+
selective_mask: Tensor,
202+
) -> Tensor:
203+
indexes = gaussian_conditional.build_indexes(scales)
204+
y_hat = means.clone()
205+
206+
for sample_index, y_string in enumerate(strings):
207+
mask = selective_mask[sample_index].reshape(-1)
208+
if not mask.any():
209+
continue
210+
indexes_i = indexes[sample_index].reshape(-1)[mask].unsqueeze(0)
211+
means_i = means[sample_index].reshape(-1)[mask].unsqueeze(0)
212+
y_hat_i = gaussian_conditional.decompress([y_string], indexes_i, means=means_i)
213+
y_hat[sample_index].reshape(-1)[mask] = y_hat_i.reshape(-1).to(y_hat.dtype)
214+
215+
return y_hat

0 commit comments

Comments
 (0)