Level 3: First positive SNR version

2026-02-15 21:16:54 +02:00
parent 4ebee28e4e
commit cd2b89bd73
28 changed files with 5043 additions and 275 deletions
@@ -0,0 +1,400 @@
+import numpy as np
+import scipy.io as sio
+import os
+
+# ------------------ LOAD LUT ------------------
+
+def load_LUT(mat_filename=None):
+    """
+    Loads the list of Huffman Codebooks (LUTs)
+
+    Returns:
+        huffLUT : list (index 1..11 used, index 0 unused)
+    """
+    if mat_filename is None:
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        mat_filename = os.path.join(current_dir, "huffCodebooks.mat")
+
+    mat = sio.loadmat(mat_filename)
+
+
+    huffCodebooks_raw = mat['huffCodebooks'].squeeze()
+
+    huffCodebooks = []
+    for i in range(11):
+        huffCodebooks.append(np.array(huffCodebooks_raw[i]))
+
+    # Build inverse VLC tables
+    invTable = [None] * 11
+
+    for i in range(11):
+        h = huffCodebooks[i][:, 2].astype(int)       # column 3
+        hlength = huffCodebooks[i][:, 1].astype(int) # column 2
+
+        hbin = []
+        for j in range(len(h)):
+            hbin.append(format(h[j], f'0{hlength[j]}b'))
+
+        invTable[i] = vlc_table(hbin)
+
+    # Build Huffman LUT dicts
+    huffLUT = [None] * 12  # index 0 unused
+    params = [
+        (4, 1, True),
+        (4, 1, True),
+        (4, 2, False),
+        (4, 2, False),
+        (2, 4, True),
+        (2, 4, True),
+        (2, 7, False),
+        (2, 7, False),
+        (2, 12, False),
+        (2, 12, False),
+        (2, 16, False),
+    ]
+
+    for i, (nTupleSize, maxAbs, signed) in enumerate(params, start=1):
+        huffLUT[i] = {
+            'LUT': huffCodebooks[i-1],
+            'invTable': invTable[i-1],
+            'codebook': i,
+            'nTupleSize': nTupleSize,
+            'maxAbsCodeVal': maxAbs,
+            'signedValues': signed
+        }
+
+    return huffLUT
+
+def vlc_table(code_array):
+    """
+    codeArray: list of strings, each string is a Huffman codeword (e.g. '0101')
+    returns:
+        h : NumPy array of shape (num_nodes, 3)
+            columns:
+              [ next_if_0 , next_if_1 , symbol_index ]
+    """
+    h = np.zeros((1, 3), dtype=int)
+
+    for code_index, code in enumerate(code_array, start=1):
+        word = [int(bit) for bit in code]
+        h_index = 0
+
+        for bit in word:
+            k = bit
+            next_node = h[h_index, k]
+            if next_node == 0:
+                h = np.vstack([h, [0, 0, 0]])
+                new_index = h.shape[0] - 1
+                h[h_index, k] = new_index
+                h_index = new_index
+            else:
+                h_index = next_node
+
+        h[h_index, 2] = code_index
+
+    return h
+
+# ------------------ ENCODE ------------------
+
+def encode_huff(coeff_sec, huff_LUT_list, force_codebook = None):
+    """
+    Huffman-encode a sequence of quantized coefficients.
+
+    This function selects the appropriate Huffman codebook based on the
+    maximum absolute value of the input coefficients, encodes the coefficients
+    into a binary Huffman bitstream, and returns both the bitstream and the
+    selected codebook index.
+
+    This is the Python equivalent of the MATLAB `encodeHuff.m` function used
+    in audio/image coding (e.g., scale factor band encoding). The input
+    coefficient sequence is grouped into fixed-size tuples as defined by
+    the chosen Huffman LUT. Zero-padding may be applied internally.
+
+    Parameters
+    ----------
+    coeff_sec : array_like of int
+        1-D array of quantized integer coefficients to encode.
+        Typically corresponds to a "section" or scale-factor band.
+
+    huff_LUT_list : list
+        List of Huffman lookup-table dictionaries as returned by `loadLUT()`.
+        Index 1..11 correspond to valid Huffman codebooks.
+        Index 0 is unused.
+
+    Returns
+    -------
+    huffSec : str
+        Huffman-encoded bitstream represented as a string of '0' and '1'
+        characters.
+
+    huffCodebook : int
+        Index (1..11) of the Huffman codebook used for encoding.
+        A value of 0 indicates a special all-zero section.
+    """
+    if force_codebook is not None:
+            return huff_LUT_code_1(huff_LUT_list[force_codebook], coeff_sec)
+    
+    maxAbsVal = np.max(np.abs(coeff_sec))
+
+    if maxAbsVal == 0:
+        huffCodebook = 0
+        huffSec = huff_LUT_code_0()
+
+    elif maxAbsVal == 1:
+        candidates = [1, 2]
+        huffSec1 = huff_LUT_code_1(huff_LUT_list[candidates[0]], coeff_sec)
+        huffSec2 = huff_LUT_code_1(huff_LUT_list[candidates[1]], coeff_sec)
+        if len(huffSec1) <= len(huffSec2):
+            huffSec = huffSec1
+            huffCodebook = candidates[0]
+        else:
+            huffSec = huffSec2
+            huffCodebook = candidates[1]
+
+    elif maxAbsVal == 2:
+        candidates = [3, 4]
+        huffSec1 = huff_LUT_code_1(huff_LUT_list[candidates[0]], coeff_sec)
+        huffSec2 = huff_LUT_code_1(huff_LUT_list[candidates[1]], coeff_sec)
+        if len(huffSec1) <= len(huffSec2):
+            huffSec = huffSec1
+            huffCodebook = candidates[0]
+        else:
+            huffSec = huffSec2
+            huffCodebook = candidates[1]
+
+    elif maxAbsVal in (3, 4):
+        candidates = [5, 6]
+        huffSec1 = huff_LUT_code_1(huff_LUT_list[candidates[0]], coeff_sec)
+        huffSec2 = huff_LUT_code_1(huff_LUT_list[candidates[1]], coeff_sec)
+        if len(huffSec1) <= len(huffSec2):
+            huffSec = huffSec1
+            huffCodebook = candidates[0]
+        else:
+            huffSec = huffSec2
+            huffCodebook = candidates[1]
+
+    elif maxAbsVal in (5, 6, 7):
+        candidates = [7, 8]
+        huffSec1 = huff_LUT_code_1(huff_LUT_list[candidates[0]], coeff_sec)
+        huffSec2 = huff_LUT_code_1(huff_LUT_list[candidates[1]], coeff_sec)
+        if len(huffSec1) <= len(huffSec2):
+            huffSec = huffSec1
+            huffCodebook = candidates[0]
+        else:
+            huffSec = huffSec2
+            huffCodebook = candidates[1]
+
+    elif maxAbsVal in (8, 9, 10, 11, 12):
+        candidates = [9, 10]
+        huffSec1 = huff_LUT_code_1(huff_LUT_list[candidates[0]], coeff_sec)
+        huffSec2 = huff_LUT_code_1(huff_LUT_list[candidates[1]], coeff_sec)
+        if len(huffSec1) <= len(huffSec2):
+            huffSec = huffSec1
+            huffCodebook = candidates[0]
+        else:
+            huffSec = huffSec2
+            huffCodebook = candidates[1]
+
+    elif maxAbsVal in (13, 14, 15):
+        huffCodebook = 11
+        huffSec = huff_LUT_code_1(huff_LUT_list[huffCodebook], coeff_sec)
+
+    else:
+        huffCodebook = 11
+        huffSec = huff_LUT_code_ESC(huff_LUT_list[huffCodebook], coeff_sec)
+
+    return huffSec, huffCodebook
+
+def huff_LUT_code_1(huff_LUT, coeff_sec):
+    LUT = huff_LUT['LUT']
+    nTupleSize = huff_LUT['nTupleSize']
+    maxAbsCodeVal = huff_LUT['maxAbsCodeVal']
+    signedValues = huff_LUT['signedValues']
+
+    numTuples = int(np.ceil(len(coeff_sec) / nTupleSize))
+
+    if signedValues:
+        coeff = coeff_sec + maxAbsCodeVal
+        base = 2 * maxAbsCodeVal + 1
+    else:
+        coeff = coeff_sec
+        base = maxAbsCodeVal + 1
+
+    coeffPad = np.zeros(numTuples * nTupleSize, dtype=int)
+    coeffPad[:len(coeff)] = coeff
+
+    huffSec = []
+
+    powers = base ** np.arange(nTupleSize - 1, -1, -1)
+
+    for i in range(numTuples):
+        nTuple = coeffPad[i*nTupleSize:(i+1)*nTupleSize]
+        huffIndex = int(np.abs(nTuple) @ powers)
+
+        hexVal = LUT[huffIndex, 2]
+        huffLen = LUT[huffIndex, 1]
+
+        bits = format(int(hexVal), f'0{int(huffLen)}b')
+
+        if signedValues:
+            huffSec.append(bits)
+        else:
+            signBits = ''.join('1' if v < 0 else '0' for v in nTuple)
+            huffSec.append(bits + signBits)
+
+    return ''.join(huffSec)
+
+def huff_LUT_code_0():
+    return ''
+
+def huff_LUT_code_ESC(huff_LUT, coeff_sec):
+    LUT = huff_LUT['LUT']
+    nTupleSize = huff_LUT['nTupleSize']
+    maxAbsCodeVal = huff_LUT['maxAbsCodeVal']
+
+    numTuples = int(np.ceil(len(coeff_sec) / nTupleSize))
+    base = maxAbsCodeVal + 1
+
+    coeffPad = np.zeros(numTuples * nTupleSize, dtype=int)
+    coeffPad[:len(coeff_sec)] = coeff_sec
+
+    huffSec = []
+    powers = base ** np.arange(nTupleSize - 1, -1, -1)
+
+    for i in range(numTuples):
+        nTuple = coeffPad[i*nTupleSize:(i+1)*nTupleSize]
+
+        lnTuple = nTuple.astype(float)
+        lnTuple[lnTuple == 0] = np.finfo(float).eps
+
+        N4 = np.maximum(0, np.floor(np.log2(np.abs(lnTuple))).astype(int))
+        N = np.maximum(0, N4 - 4)
+        esc = np.abs(nTuple) > 15
+
+        nTupleESC = nTuple.copy()
+        nTupleESC[esc] = np.sign(nTupleESC[esc]) * 16
+
+        huffIndex = int(np.abs(nTupleESC) @ powers)
+
+        hexVal = LUT[huffIndex, 2]
+        huffLen = LUT[huffIndex, 1]
+
+        bits = format(int(hexVal), f'0{int(huffLen)}b')
+
+        escSeq = ''
+        for k in range(nTupleSize):
+            if esc[k]:
+                escSeq += '1' * N[k]
+                escSeq += '0'
+                escSeq += format(abs(nTuple[k]) - (1 << N4[k]), f'0{N4[k]}b')
+
+        signBits = ''.join('1' if v < 0 else '0' for v in nTuple)
+        huffSec.append(bits + signBits + escSeq)
+
+    return ''.join(huffSec)
+
+# ------------------ DECODE ------------------
+
+def decode_huff(huff_sec, huff_LUT):
+    """
+    Decode a Huffman-encoded stream.
+
+    Parameters
+    ----------
+    huff_sec : array-like of int or str
+        Huffman encoded stream as a sequence of 0 and 1 (string or list/array).
+    huff_LUT : dict
+        Huffman lookup table with keys:
+            - 'invTable': inverse table (numpy array)
+            - 'codebook': codebook number
+            - 'nTupleSize': tuple size
+            - 'maxAbsCodeVal': maximum absolute code value
+            - 'signedValues': True/False
+
+    Returns
+    -------
+    decCoeffs : list of int
+        Decoded quantized coefficients.
+    """
+
+    h = huff_LUT['invTable']
+    huffCodebook = huff_LUT['codebook']
+    nTupleSize = huff_LUT['nTupleSize']
+    maxAbsCodeVal = huff_LUT['maxAbsCodeVal']
+    signedValues = huff_LUT['signedValues']
+
+    # Convert string to array of ints
+    if isinstance(huff_sec, str):
+        huff_sec = np.array([int(b) for b in huff_sec])
+
+    eos = False
+    decCoeffs = []
+    streamIndex = 0
+
+    while not eos:
+        wordbit = 0
+        r = 0  # start at root
+
+        # Decode Huffman word using inverse table
+        while True:
+            b = huff_sec[streamIndex + wordbit]
+            wordbit += 1
+            rOld = r
+            r = h[rOld, b]
+            if h[r, 0] == 0 and h[r, 1] == 0:
+                symbolIndex = h[r, 2] - 1  # zero-based
+                streamIndex += wordbit
+                break
+
+        # Decode n-tuple magnitudes
+        if signedValues:
+            base = 2 * maxAbsCodeVal + 1
+            nTupleDec = []
+            tmp = symbolIndex
+            for p in reversed(range(nTupleSize)):
+                val = tmp // (base ** p)
+                nTupleDec.append(val - maxAbsCodeVal)
+                tmp = tmp % (base ** p)
+            nTupleDec = np.array(nTupleDec)
+        else:
+            base = maxAbsCodeVal + 1
+            nTupleDec = []
+            tmp = symbolIndex
+            for p in reversed(range(nTupleSize)):
+                val = tmp // (base ** p)
+                nTupleDec.append(val)
+                tmp = tmp % (base ** p)
+            nTupleDec = np.array(nTupleDec)
+
+            # Apply sign bits
+            nTupleSignBits = huff_sec[streamIndex:streamIndex + nTupleSize]
+            nTupleSign = -(np.sign(nTupleSignBits - 0.5))
+            streamIndex += nTupleSize
+            nTupleDec = nTupleDec * nTupleSign
+
+        # Handle escape sequences
+        escIndex = np.where(np.abs(nTupleDec) == 16)[0]
+        if huffCodebook == 11 and escIndex.size > 0:
+            for idx in escIndex:
+                N = 0
+                b = huff_sec[streamIndex]
+                while b:
+                    N += 1
+                    b = huff_sec[streamIndex + N]
+                streamIndex += N +1
+                N4 = N + 4
+                escape_word = huff_sec[streamIndex:streamIndex + N4]
+                escape_value = 2 ** N4 + int("".join(map(str, escape_word)), 2)
+                nTupleDec[idx] = escape_value
+                streamIndex += N4
+            # Apply signs again
+            nTupleDec[escIndex] *= nTupleSign[escIndex]
+
+        decCoeffs.extend(nTupleDec.tolist())
+
+        if streamIndex >= len(huff_sec):
+            eos = True
+
+    return decCoeffs
+    
+