kwz adpcm decoder

kwz_adpcm_decoder.py

import numpy as np 
import wave
import audioop
from sys import argv

step_table = np.array([
  7,     8,     9,     10,    11,    12,    13,    14,    16,    17,
  19,    21,    23,    25,    28,    31,    34,    37,    41,    45,
  50,    55,    60,    66,    73,    80,    88,    97,    107,   118,
  130,   143,   157,   173,   190,   209,   230,   253,   279,   307,
  337,   371,   408,   449,   494,   544,   598,   658,   724,   796,
  876,   963,   1060,  1166,  1282,  1411,  1552,  1707,  1878,  2066,
  2272,  2499,  2749,  3024,  3327,  3660,  4026,  4428,  4871,  5358,
  5894,  6484,  7132,  7845,  8630,  9493,  10442, 11487, 12635, 13899,
  15289, 16818, 18500, 20350, 22385, 24623, 27086, 29794, 32767, 0
], dtype=np.int16)

index_table_2bit = np.array([
  -1,  2,
  -1,  2,
], dtype=np.int8)

index_table_4bit = np.array([
  -1, -1, -1, -1, 2, 4, 6, 8,
  -1, -1, -1, -1, 2, 4, 6, 8,
], dtype=np.int8)

diff_table_2bit = np.zeros(90 * 4, dtype=np.int16)

for sample in range (4):
  for step_index in range(90):
    step = step_table[step_index]
    diff = step >> 3
    if (sample & 1): diff += step;
    if (sample & 2): diff = -diff;
    diff_table_2bit[(sample + 4 * step_index)] = diff

diff_table_4bit = np.zeros(90 * 16, dtype=np.int16)

for sample in range (16):
  for step_index in range(90):
    step = step_table[step_index]
    diff = step >> 3;
    if (sample & 4): diff += step;
    if (sample & 2): diff += step >> 1;
    if (sample & 1): diff += step >> 2;
    if (sample & 8): diff = -diff;
    diff_table_4bit[(sample + 16 * step_index)] = diff


with open(argv[1], "rb") as adpcm:
  # pcm output buffer - long enough for 60 seconds of audio at 16364 hz
  pcm = np.zeros(16364 * 60, dtype="<u2");
  pcmOffset = 0

  # initial decoder state
  prev_diff = 0
  prev_step_index = 40

  for byte in adpcm.read():

    bit_pos = 0
    while (bit_pos < 8):

      # if we can only read a 2-bit value or if the previous step index is < 0x12
      if (prev_step_index < 0x12 or bit_pos == 6):
        # isolate 2-bit sample
        sample = (byte >> bit_pos) & 0x3
        # get diff
        diff = prev_diff + diff_table_2bit[sample + 4 * prev_step_index]
        # get step index
        step_index = prev_step_index + index_table_2bit[sample]
        bit_pos += 2

      else:
        # isolate 4-bit sample
        sample = (byte >> bit_pos) & 0xF
        # get diff
        diff = prev_diff + diff_table_4bit[sample + 16 * prev_step_index]
        # get step index
        step_index = prev_step_index + index_table_4bit[sample]
        bit_pos += 4

      # clamp step index and diff
      step_index = max(0, min(step_index, 79))
      diff = max(-2048, min(diff, 2048))

      # add result to output buffer
      pcm[pcmOffset] = diff * 16
      pcmOffset += 1

      # set prev decoder state
      prev_step_index = step_index
      prev_diff = diff

  # output as .wav
  audio = wave.open(argv[2], "wb")
  audio.setnchannels(1)
  audio.setsampwidth(2)
  audio.setframerate(16364)
  audio.writeframes(pcm[:pcmOffset].tobytes())
  audio.close()