ArithmeticDecoder, see `ArithmeticCoder` for a detailed explanation. Note that this must be called with **exactly** the same parameters and sequence of quantized cdf as the arithmetic encoder or the wrong values will be decoded. If the AC encoder current range is [L, H], with `L` and `
| 168 | |
| 169 | |
| 170 | class ArithmeticDecoder: |
| 171 | """ArithmeticDecoder, see `ArithmeticCoder` for a detailed explanation. |
| 172 | |
| 173 | Note that this must be called with **exactly** the same parameters and sequence |
| 174 | of quantized cdf as the arithmetic encoder or the wrong values will be decoded. |
| 175 | |
| 176 | If the AC encoder current range is [L, H], with `L` and `H` having the some common |
| 177 | prefix (i.e. the same most significant bits), then this prefix will be flushed to the stream. |
| 178 | For instances, having read 3 bits `b1 b2 b3`, we know that `[L, H]` is contained inside |
| 179 | `[b1 b2 b3 0 ... 0 b1 b3 b3 1 ... 1]`. Now this specific sub-range can only be obtained |
| 180 | for a specific sequence of symbols and a binary-search allows us to decode those symbols. |
| 181 | At some point, the prefix `b1 b2 b3` will no longer be sufficient to decode new symbols, |
| 182 | and we will need to read new bits from the stream and repeat the process. |
| 183 | |
| 184 | """ |
| 185 | def __init__(self, fo: tp.IO[bytes], total_range_bits: int = 24): |
| 186 | self.total_range_bits = total_range_bits |
| 187 | self.low: int = 0 |
| 188 | self.high: int = 0 |
| 189 | self.current: int = 0 |
| 190 | self.max_bit: int = -1 |
| 191 | self.unpacker = BitUnpacker(bits=1, fo=fo) # we pull single bits at a time. |
| 192 | # Following is for debugging |
| 193 | self._dbg: tp.List[tp.Any] = [] |
| 194 | self._dbg2: tp.List[tp.Any] = [] |
| 195 | self._last: tp.Any = None |
| 196 | |
| 197 | @property |
| 198 | def delta(self) -> int: |
| 199 | return self.high - self.low + 1 |
| 200 | |
| 201 | def _flush_common_prefix(self): |
| 202 | # Given the current range [L, H], if both have a common prefix, |
| 203 | # we know we can remove it from our representation to avoid handling large numbers. |
| 204 | while self.max_bit >= 0: |
| 205 | b1 = self.low >> self.max_bit |
| 206 | b2 = self.high >> self.max_bit |
| 207 | if b1 == b2: |
| 208 | self.low -= (b1 << self.max_bit) |
| 209 | self.high -= (b1 << self.max_bit) |
| 210 | self.current -= (b1 << self.max_bit) |
| 211 | assert self.high >= self.low |
| 212 | assert self.low >= 0 |
| 213 | self.max_bit -= 1 |
| 214 | else: |
| 215 | break |
| 216 | |
| 217 | def pull(self, quantized_cdf: torch.Tensor) -> tp.Optional[int]: |
| 218 | """Pull a symbol, reading as many bits from the stream as required. |
| 219 | This returns `None` when the stream has been exhausted. |
| 220 | |
| 221 | Args: |
| 222 | quantized_cdf (torch.Tensor): use `build_stable_quantized_cdf` |
| 223 | to build this from your pdf estimate. This must be **exatly** |
| 224 | the same cdf as the one used at encoding time. |
| 225 | """ |
| 226 | while self.delta < 2 ** self.total_range_bits: |
| 227 | bit = self.unpacker.pull() |