diff --git a/coder/bit_alloc.h b/coder/bit_alloc.h index 39fb056..8f2a249 100644 --- a/coder/bit_alloc.h +++ b/coder/bit_alloc.h @@ -65,7 +65,7 @@ inline DualAllocResult allocate_bits_dual( if (stereo) { if (target_kbps < 128) { - if (target_kbps >= 64) { + if (target_kbps < 96) { if (energy0[0] > energy1[0]) { bits0[0] = 4; bits1[0] = 3; @@ -74,10 +74,10 @@ inline DualAllocResult allocate_bits_dual( bits1[0] = 4; } } else { - if (target_kbps >= 0) { - bits0[0] = 4; - bits1[0] = 3; - } + // if (target_kbps >= 0) { + bits0[0] = 4; + bits1[0] = 4; + // } // bits0[0] = 4; // bits1[0] = 3; // if (energy0[0] > energy1[0]) { diff --git a/coder/codec_utils.h b/coder/codec_utils.h index ef8c0ef..c33697c 100644 --- a/coder/codec_utils.h +++ b/coder/codec_utils.h @@ -6,6 +6,7 @@ #include inline int get_scale_bits(int band_idx) { + return 8; if (band_idx < 1) return 12; return 8; } diff --git a/coder/coder.cpp b/coder/coder.cpp index 891cdaa..437bfc8 100644 --- a/coder/coder.cpp +++ b/coder/coder.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #define DR_WAV_IMPLEMENTATION #include "dr_wav.h" #include "wavelet.h" @@ -27,38 +28,40 @@ #include "energy_shape.h" #include "sbr_encode.h" #include "pow_filter.h" + #ifndef M_PI #define M_PI 3.14159265358979323846 #endif -static int get_adaptive_is_max_base(float target_kbps, const EncoderSettings& s) { +float SBR_W0 = 1.0f; +float SBR_W1 = 9.0f; + +static int get_adaptive_is_max_base(float target_kbps, const EncoderSettings& s) +{ int base_default = get_is_base_original(target_kbps); if (s.adaptive_is_base_override > 0) return std::clamp(s.adaptive_is_base_override, base_default, 16); if (s.adaptive_is_base_override == 0) return base_default; - if (target_kbps < s.adaptive_is_kbps_min || - target_kbps > s.adaptive_is_kbps_max) + if (target_kbps < s.adaptive_is_kbps_min || target_kbps > s.adaptive_is_kbps_max) return base_default; - + float k0, k1; - int b0, b1; + int b0, b1; if (target_kbps <= 64.0f) { k0 = k1 = target_kbps; b0 = b1 = s.adaptive_is_base_at_64; } else if (target_kbps <= 96.0f) { - k0 = 64.0f; k1 = 96.0f; - b0 = s.adaptive_is_base_at_64; - b1 = s.adaptive_is_base_at_96; + k0 = 64.0f; k1 = 96.0f; + b0 = s.adaptive_is_base_at_64; b1 = s.adaptive_is_base_at_96; } else if (target_kbps <= 128.0f) { - k0 = 96.0f; k1 = 128.0f; - b0 = s.adaptive_is_base_at_96; - b1 = s.adaptive_is_base_at_128; + k0 = 96.0f; k1 = 128.0f; + b0 = s.adaptive_is_base_at_96; b1 = s.adaptive_is_base_at_128; } else { k0 = 128.0f; k1 = 160.0f; - b0 = s.adaptive_is_base_at_128; - b1 = s.adaptive_is_base_at_160; + b0 = s.adaptive_is_base_at_128; b1 = s.adaptive_is_base_at_160; } + float t = 0.0f; if (std::abs(k1 - k0) > 1e-6f) t = (target_kbps - k0) / (k1 - k0); @@ -69,10 +72,18 @@ static int get_adaptive_is_max_base(float target_kbps, const EncoderSettings& s) } std::vector compress_block_adaptive_joint( - const std::vector& block, PWPT& wpt, int num_channels, int level, - int target_bits_budget, float target_kbps, int sr, int block_size, - bool enable_ms, bool block_is_full, - uint32_t block_index, PredContext& pred_ctx) + const std::vector& block, + PWPT& wpt, + int num_channels, + int level, + int target_bits_budget, + float target_kbps, + int sr, + int block_size, + bool enable_ms, + bool block_is_full, + uint32_t block_index, + PredContext& pred_ctx) { const float eps = 1e-12f; const bool stereo = (num_channels == 2); @@ -84,10 +95,10 @@ std::vector compress_block_adaptive_joint( size_t spc = block.size() / 2; std::vector left(spc), right(spc); for (size_t i = 0; i < spc; ++i) { - left[i] = block[i * 2]; + left[i] = block[i * 2]; right[i] = block[i * 2 + 1]; } - left_coeffs = wpt.wpt(left, level, target_kbps, sr, num_channels); + left_coeffs = wpt.wpt(left, level, target_kbps, sr, num_channels); right_coeffs = wpt.wpt(right, level, target_kbps, sr, num_channels); coeff_counts.resize(left_coeffs.size()); for (size_t i = 0; i < left_coeffs.size(); ++i) @@ -99,9 +110,8 @@ std::vector compress_block_adaptive_joint( coeff_counts[i] = (int)left_coeffs[i].size(); } - int band_count = (int)coeff_counts.size(); + int band_count = (int)coeff_counts.size(); int total_bands = 1 << level; - ModeState& pred_state = (level == 5) ? pred_ctx.long_state : pred_ctx.short_state; pred_state.ensure(band_count); @@ -122,10 +132,11 @@ std::vector compress_block_adaptive_joint( ch0_original[i] = left_coeffs[i]; int is_start_default = band_count; - int is_start_max = band_count; + int is_start_max = band_count; std::vector> is_r_vals(band_count); for (auto& a : is_r_vals) a.fill(0.5f); - std::vector is_inv_flags(band_count, false); + std::vector> is_inv_flags(band_count); + for (auto& a : is_inv_flags) a.fill(false); std::vector is_use_segmented(band_count, false); if (use_is) { @@ -136,21 +147,16 @@ std::vector compress_block_adaptive_joint( bool adaptive_is = false; if (use_is) { - bool in_range = - (target_kbps >= SETTINGS.adaptive_is_kbps_min && - target_kbps <= SETTINGS.adaptive_is_kbps_max); - bool allow_adaptive = - SETTINGS.enable_adaptive_is && - (SETTINGS.adaptive_is_base_override != 0) && - (in_range || SETTINGS.adaptive_is_base_override > 0); - + bool in_range = (target_kbps >= SETTINGS.adaptive_is_kbps_min && + target_kbps <= SETTINGS.adaptive_is_kbps_max); + bool allow_adaptive = SETTINGS.enable_adaptive_is && + (SETTINGS.adaptive_is_base_override != 0) && + (in_range || SETTINGS.adaptive_is_base_override > 0); if (allow_adaptive) { int base_max = get_adaptive_is_max_base(target_kbps, SETTINGS); is_start_max = get_is_start_from_base(base_max, total_bands); - if (is_start_max < is_start_default) - is_start_max = is_start_default; - if (is_start_max > band_count) - is_start_max = band_count; + if (is_start_max < is_start_default) is_start_max = is_start_default; + if (is_start_max > band_count) is_start_max = band_count; adaptive_is = true; } } @@ -158,37 +164,70 @@ std::vector compress_block_adaptive_joint( if (stereo) { int low_end = use_is ? is_start_default : band_count; for (int i = 0; i < low_end; ++i) { - const auto& left_band = left_coeffs[i]; - const auto& right_band = right_coeffs[i]; - int n = (int)left_band.size(); + const auto& lb = left_coeffs[i]; + const auto& rb = right_coeffs[i]; + int n = (int)lb.size(); if (n == 0) { - ch0_bands[i] = left_band; - ch1_bands[i] = right_band; + ch0_bands[i] = lb; + ch1_bands[i] = rb; continue; } - - float El = 0.0f, Er = 0.0f; + float El = 0, Er = 0, Em = 0, Es = 0; for (int j = 0; j < n; ++j) { - El += left_band[j] * left_band[j]; - Er += right_band[j] * right_band[j]; + El += lb[j] * lb[j]; + Er += rb[j] * rb[j]; } - - auto [mid, side] = mid_side(left_band, right_band); - float Em = 0.0f, Es = 0.0f; + auto [mid, side] = mid_side(lb, rb); for (int j = 0; j < n; ++j) { - Em += mid[j] * mid[j]; + Em += mid[j] * mid[j]; Es += side[j] * side[j]; } - - bool use_ms = use_mid_side(El, Er, Em, Es, enable_ms, target_kbps); - if (use_ms) { - mode_ms[i] = 1; + if (use_mid_side(El, Er, Em, Es, enable_ms, target_kbps)) { + mode_ms[i] = 1; ch0_bands[i] = std::move(mid); ch1_bands[i] = std::move(side); } else { - mode_ms[i] = 0; - ch0_bands[i] = left_band; - ch1_bands[i] = right_band; + mode_ms[i] = 0; + ch0_bands[i] = lb; + ch1_bands[i] = rb; + } + } + + if (use_is) { + BandEnergy be = compute_band_energy(ch0_bands, ch1_bands, is_band, stereo, band_count); + int is_seg_threshold = 8 * total_bands / 16; + for (int i = is_start_default; i < band_count; ++i) { + bool force_is = !adaptive_is || (i >= is_start_max); + bool choose_is = force_is; + if (!force_is) { + choose_is = should_use_is_band( + left_coeffs[i], right_coeffs[i], enable_ms, target_kbps, + SETTINGS.is_strong_ratio_threshold, + SETTINGS.is_equal_ratio_threshold, + SETTINGS.is_equal_ratio_threshold_low_bitrate); + } + if (choose_is) { + is_band[i] = 1; + mode_ms[i] = 1; + std::array r_arr; + std::array inv_arr; + bool use_segmented; + std::vector Y; + compute_is_parameters_ex( + left_coeffs[i], right_coeffs[i], Y, r_arr, + inv_arr, use_segmented, i, total_bands); + is_r_vals[i] = r_arr; + is_inv_flags[i] = inv_arr; + is_use_segmented[i] = (i < is_seg_threshold); + ch0_bands[i] = std::move(Y); + ch1_bands[i] = ch0_bands[i]; + } else { + is_band[i] = 0; + mode_ms[i] = 0; + auto [mid, side] = mid_side(left_coeffs[i], right_coeffs[i]); + ch0_bands[i] = mid; + ch1_bands[i] = side; + } } } @@ -249,125 +288,87 @@ std::vector compress_block_adaptive_joint( std::vector energy1 = std::move(be.e1); apply_base_pow(energy0, energy1, band_count); - bool has_transient = detect_strong_transient( - energy0, energy1, band_count, stereo, - SETTINGS.transient_ratio_threshold); - - apply_adaptive_pow(energy0, energy1, band_count, stereo, - target_kbps, has_transient); + energy0, energy1, band_count, stereo, SETTINGS.transient_ratio_threshold); + apply_adaptive_pow(energy0, energy1, band_count, stereo, target_kbps, has_transient); apply_narrowband_correction(energy0, energy1, band_count, stereo); - apply_highfreq_attenuation(energy0, energy1, band_count, - target_kbps, num_channels, has_transient); + apply_highfreq_attenuation( + energy0, energy1, band_count, target_kbps, num_channels, has_transient); static std::vector> prev_ch0, prev_ch1; std::vector priority0, priority1; if (stereo) { - priority0 = compute_channel_priority(ch0_bands, prev_ch0, coeff_counts, - target_kbps, sr, level, false); - priority1 = compute_channel_priority(ch1_bands, prev_ch1, coeff_counts, - target_kbps, sr, level, true); + priority0 = compute_channel_priority( + ch0_bands, prev_ch0, coeff_counts, target_kbps, sr, level, false); + priority1 = compute_channel_priority( + ch1_bands, prev_ch1, coeff_counts, target_kbps, sr, level, true); } else { - priority0 = compute_channel_priority(ch0_bands, prev_ch0, coeff_counts, - target_kbps, sr, level, false); + priority0 = compute_channel_priority( + ch0_bands, prev_ch0, coeff_counts, target_kbps, sr, level, false); priority1.assign(band_count, 0.0f); } std::vector min_bits(band_count, 0); std::vector max_bits(band_count, 10); - if (use_is) { for (int i = 0; i < band_count; ++i) { if (is_band[i]) { - min_bits[i] = 0; - priority1[i] = 0.0f; - energy1[i] = 0.0f; + min_bits[i] = 0; + priority1[i] = 0.0f; + energy1[i] = 0.0f; } } } - int bits_header = 0; - if (stereo) bits_header += ((band_count + 7) / 8) * 8; - bits_header += ((band_count + 7) / 8) * 8; - if (stereo) bits_header += ((band_count + 7) / 8) * 8; - bits_header += 8; - - int header_bytes_est = (bits_header + 7) / 8 + 2; - int payload_budget = target_bits_budget - header_bytes_est * 8; - if (payload_budget < 0) - payload_budget = target_bits_budget / 2; - + int header_bytes_est = 3; + int payload_budget = target_bits_budget - header_bytes_est * 8; + if (payload_budget < 0) payload_budget = target_bits_budget / 2; int reservoir_max = (int)(payload_budget * SETTINGS.reservoir_max_factor); DualAllocResult alloc = allocate_bits_dual( priority0, priority1, coeff_counts, payload_budget, - min_bits, max_bits, energy0, energy1, - 0, reservoir_max, target_kbps, num_channels == 2); + min_bits, max_bits, energy0, energy1, 0, reservoir_max, + target_kbps, num_channels == 2); - std::vector bits0 = alloc.bits0; - std::vector bits1 = alloc.bits1; + std::vector bits0 = std::move(alloc.bits0); + std::vector bits1 = std::move(alloc.bits1); std::vector active0(band_count), active1(band_count); for (int i = 0; i < band_count; ++i) { - if (stereo && is_band[i]) - bits1[i] = 0; + if (stereo && is_band[i]) bits1[i] = 0; active0[i] = (bits0[i] > 0); - if (stereo) - active1[i] = (!is_band[i] && bits1[i] > 0); - else - active1[i] = false; + if (stereo) active1[i] = (!is_band[i] && bits1[i] > 0); + else active1[i] = false; } int sbr_end = block_is_full ? (3 * band_count / 4) : 0; - SBREncodeResult sbr = sbr_analyze(ch0_original, active0, sbr_end, band_count); + SBREncodeResult sbr = sbr_analyze( + ch0_original, active0, sbr_end, band_count, SBR_W0, SBR_W1); std::vector header; - int mask_bytes = (band_count + 7) / 8; - - std::vector mode_bytes_vec((band_count + 7) / 8, 0); - for (int i = 0; i < band_count; ++i) - if (mode_ms[i]) mode_bytes_vec[i / 8] |= (1 << (i % 8)); - std::vector mask0(mask_bytes, 0); - for (int i = 0; i < band_count; ++i) - if (active0[i]) mask0[i / 8] |= (1 << (i % 8)); - std::vector mask1; - if (stereo) { - mask1.assign(mask_bytes, 0); - for (int i = 0; i < band_count; ++i) - if (active1[i]) mask1[i / 8] |= (1 << (i % 8)); - } - - header.insert(header.end(), mode_bytes_vec.begin(), mode_bytes_vec.end()); - header.insert(header.end(), mask0.begin(), mask0.end()); - if (stereo) - header.insert(header.end(), mask1.begin(), mask1.end()); - - std::vector>> for_quant0 = { ch0_bands }; + std::vector>> for_quant0(1); + for_quant0[0] = std::move(ch0_bands); QuantResult qres0 = quantize_levels(for_quant0, bits0, target_kbps); + QuantResult qres1; if (stereo) { - std::vector>> for_quant1 = { ch1_bands }; + std::vector>> for_quant1(1); + for_quant1[0] = std::move(ch1_bands); qres1 = quantize_levels(for_quant1, bits1, target_kbps); } - std::vector scale_raw0, scale_raw1; - std::vector scale_pred0, scale_pred1; - std::vector full_scale0(band_count, 0); - std::vector full_scale1(band_count, 0); - + std::vector scale_raw0, scale_raw1, scale_pred0, scale_pred1; + std::vector full_scale0(band_count, 0), full_scale1(band_count, 0); for (int i = 0; i < band_count; ++i) { if (active0[i]) { - int sb = get_scale_bits(i); + int sb = get_scale_bits(i); float step = scale_and_bits_to_step(qres0.scales[i], bits0[i]); uint32_t idx = get_scale_idx(step, sb); full_scale0[i] = idx; scale_raw0.push_back(idx); - if (can_state_pred) { uint32_t val = idx; - if (i < (int)pred_state.prev_active0.size() && - pred_state.prev_active0[i]) - { + if (i < (int)pred_state.prev_active0.size() && pred_state.prev_active0[i]) { int32_t d = (int32_t)idx - (int32_t)pred_state.prev_scale0[i]; val = zigzag_encode(d); } @@ -375,17 +376,14 @@ std::vector compress_block_adaptive_joint( } } if (stereo && active1[i]) { - int sb = get_scale_bits(i); + int sb = get_scale_bits(i); float step = scale_and_bits_to_step(qres1.scales[i], bits1[i]); uint32_t idx = get_scale_idx(step, sb); full_scale1[i] = idx; scale_raw1.push_back(idx); - if (can_state_pred) { uint32_t val = idx; - if (i < (int)pred_state.prev_active1.size() && - pred_state.prev_active1[i]) - { + if (i < (int)pred_state.prev_active1.size() && pred_state.prev_active1[i]) { int32_t d = (int32_t)idx - (int32_t)pred_state.prev_scale1[i]; val = zigzag_encode(d); } @@ -394,76 +392,116 @@ std::vector compress_block_adaptive_joint( } } - int k_scale0_raw = compute_optimal_rice_k(scale_raw0); - int k_scale1_raw = stereo ? compute_optimal_rice_k(scale_raw1) : 0; - size_t raw_scale_bits = - estimate_rice_bits(scale_raw0, k_scale0_raw) + - (stereo ? estimate_rice_bits(scale_raw1, k_scale1_raw) : 0); + int k_scale0_raw = compute_optimal_rice_k(scale_raw0); + int k_scale1_raw = stereo ? compute_optimal_rice_k(scale_raw1) : 0; + size_t raw_scale_bits = estimate_rice_bits(scale_raw0, k_scale0_raw) + + (stereo ? estimate_rice_bits(scale_raw1, k_scale1_raw) : 0); - int k_scale0_pred = 0, k_scale1_pred = 0; + int k_scale0_pred = 0; + int k_scale1_pred = 0; size_t pred_scale_bits = SIZE_MAX; if (can_state_pred) { k_scale0_pred = compute_optimal_rice_k(scale_pred0); k_scale1_pred = stereo ? compute_optimal_rice_k(scale_pred1) : 0; - pred_scale_bits = - estimate_rice_bits(scale_pred0, k_scale0_pred) + - (stereo ? estimate_rice_bits(scale_pred1, k_scale1_pred) : 0); + pred_scale_bits = estimate_rice_bits(scale_pred0, k_scale0_pred) + + (stereo ? estimate_rice_bits(scale_pred1, k_scale1_pred) : 0); } bool use_scale_pred = can_state_pred && (pred_scale_bits < raw_scale_bits); - int k_scale0 = use_scale_pred ? k_scale0_pred : k_scale0_raw; - int k_scale1 = use_scale_pred ? k_scale1_pred : k_scale1_raw; + int k_scale0 = use_scale_pred ? k_scale0_pred : k_scale0_raw; + int k_scale1 = use_scale_pred ? k_scale1_pred : k_scale1_raw; - auto scale_vals0 = use_scale_pred ? scale_pred0 : scale_raw0; - auto scale_vals1 = use_scale_pred ? scale_pred1 : scale_raw1; + auto scale_vals0 = use_scale_pred ? std::move(scale_pred0) : std::move(scale_raw0); + auto scale_vals1 = use_scale_pred ? std::move(scale_pred1) : std::move(scale_raw1); uint8_t k_scale_byte = (uint8_t)(k_scale0 & 0x07); - if (stereo) - k_scale_byte |= ((k_scale1 & 0x07) << 3); - if (block_is_full) - k_scale_byte |= 0x80; - if (use_scale_pred) - k_scale_byte |= 0x40; + if (stereo) k_scale_byte |= ((k_scale1 & 0x07) << 3); + if (block_is_full) k_scale_byte |= 0x80; + if (use_scale_pred) k_scale_byte |= 0x40; header.push_back(k_scale_byte); BitWriterMSB payload_writer; + uint32_t packed_mode = 0; + uint32_t packed_active0 = 0; + for (int i = 0; i < band_count; ++i) { + packed_mode = (packed_mode << 1) | mode_ms[i]; + packed_active0 = (packed_active0 << 1) | (active0[i] ? 1 : 0); + } + payload_writer.write_bits(packed_mode, band_count); + payload_writer.write_bits(packed_active0, band_count); + + if (stereo) { + int is_start = use_is ? get_is_start_band(target_kbps, total_bands) : band_count; + uint32_t packed_active1 = 0; + int non_is_count = 0; + for (int i = 0; i < band_count; ++i) { + bool is_is_band = use_is && (i >= is_start) && mode_ms[i]; + if (!is_is_band) { + packed_active1 = (packed_active1 << 1) | (active1[i] ? 1 : 0); + non_is_count++; + } + } + if (non_is_count > 0) { + payload_writer.write_bits(packed_active1, non_is_count); + } + } + for (uint32_t v : scale_vals0) rice_encode(payload_writer, std::vector{v}, k_scale0); - if (stereo) + if (stereo) { for (uint32_t v : scale_vals1) rice_encode(payload_writer, std::vector{v}, k_scale1); + } + + std::vector k_sub_vals0, k_sub_vals1; + std::vector k_sub_per_band0(band_count, 0), k_sub_per_band1(band_count, 0); + std::vector> uvals0(band_count), uvals1(band_count); + k_sub_vals0.reserve(band_count); + if (stereo) k_sub_vals1.reserve(band_count); for (int i = 0; i < band_count; ++i) { - if (!active0[i] || bits0[i] <= 0) continue; - const auto& qvals = qres0.quantized_per_level[i]; - std::vector uvals(qvals.size()); - for (size_t j = 0; j < qvals.size(); ++j) - uvals[j] = zigzag_encode(qvals[j]); - int k_sub = compute_optimal_rice_k(uvals); - payload_writer.write_bits(k_sub, 3); - rice_encode(payload_writer, uvals, k_sub); + if (active0[i] && bits0[i] > 0) { + const auto& qvals = qres0.quantized_per_level[i]; + auto& uv = uvals0[i]; + uv.resize(qvals.size()); + for (size_t j = 0; j < qvals.size(); ++j) + uv[j] = zigzag_encode(qvals[j]); + int k_sub = compute_optimal_rice_k(uv); + k_sub_per_band0[i] = k_sub; + k_sub_vals0.push_back(k_sub); + } + if (stereo && active1[i] && bits1[i] > 0) { + const auto& qvals = qres1.quantized_per_level[i]; + auto& uv = uvals1[i]; + uv.resize(qvals.size()); + for (size_t j = 0; j < qvals.size(); ++j) + uv[j] = zigzag_encode(qvals[j]); + int k_sub = compute_optimal_rice_k(uv); + k_sub_per_band1[i] = k_sub; + k_sub_vals1.push_back(k_sub); + } } + int k_sub_k0 = k_sub_vals0.empty() ? 0 : std::min(compute_optimal_rice_k(k_sub_vals0), 3); + int k_sub_k1 = 0; + if (stereo && !k_sub_vals1.empty()) + k_sub_k1 = std::min(compute_optimal_rice_k(k_sub_vals1), 3); + + payload_writer.write_bits(k_sub_k0, 2); + if (stereo) payload_writer.write_bits(k_sub_k1, 2); + rice_encode(payload_writer, k_sub_vals0, k_sub_k0); + if (stereo) rice_encode(payload_writer, k_sub_vals1, k_sub_k1); + + for (int i = 0; i < band_count; ++i) { + if (active0[i] && bits0[i] > 0) + rice_encode(payload_writer, uvals0[i], k_sub_per_band0[i]); + } if (stereo) { for (int i = 0; i < band_count; ++i) { - if (!active1[i] || bits1[i] <= 0) continue; - const auto& qvals = qres1.quantized_per_level[i]; - std::vector uvals(qvals.size()); - for (size_t j = 0; j < qvals.size(); ++j) - uvals[j] = zigzag_encode(qvals[j]); - int k_sub = compute_optimal_rice_k(uvals); - payload_writer.write_bits(k_sub, 3); - rice_encode(payload_writer, uvals, k_sub); + if (active1[i] && bits1[i] > 0) + rice_encode(payload_writer, uvals1[i], k_sub_per_band1[i]); } } - payload_writer.flush(); - auto payload = payload_writer.data(); - uint16_t payload_len = (uint16_t)payload.size(); - header.push_back(payload_len & 0xFF); - header.push_back((payload_len >> 8) & 0xFF); - - std::vector out = header; - out.insert(out.end(), payload.begin(), payload.end()); if (use_is) { std::vector is_indices; @@ -471,29 +509,48 @@ std::vector compress_block_adaptive_joint( if (is_band[i]) is_indices.push_back(i); if (!is_indices.empty()) { - BitWriterMSB r_writer; + int segmented_threshold = 8 * total_bands / 16; + int total_vals = 0; + for (int idx : is_indices) + total_vals += (idx < segmented_threshold) ? 4 : 1; + + uint32_t current_chunk = 0; + int bits_in_chunk = 0; + for (int i : is_indices) { + bool is_seg = (i < segmented_threshold); + int bits_to_add = is_seg ? 4 : 1; + for (int s = 0; s < bits_to_add; ++s) { + int val = is_seg ? (is_inv_flags[i][s] ? 1 : 0) : (is_inv_flags[i][0] ? 1 : 0); + current_chunk = (current_chunk << 1) | val; + bits_in_chunk++; + if (bits_in_chunk == 32) { + payload_writer.write_bits(current_chunk, 32); + current_chunk = 0; + bits_in_chunk = 0; + } + } + } + if (bits_in_chunk > 0) { + payload_writer.write_bits(current_chunk, bits_in_chunk); + } + + std::vector is_zz_vals; + is_zz_vals.reserve(total_vals); for (int i : is_indices) { int bits = get_r_bits(i); if (is_use_segmented[i]) { for (int s = 0; s < 4; ++s) { - uint32_t q = quantize_r(is_r_vals[i][s], bits); - r_writer.write_bits(q, bits); + int32_t q = quantize_r_centered(is_r_vals[i][s], bits); + is_zz_vals.push_back(zigzag_encode(q)); } } else { - uint32_t q = quantize_r(is_r_vals[i][0], bits); - r_writer.write_bits(q, bits); + int32_t q = quantize_r_centered(is_r_vals[i][0], bits); + is_zz_vals.push_back(zigzag_encode(q)); } } - r_writer.flush(); - auto r_bytes = r_writer.data(); - out.insert(out.end(), r_bytes.begin(), r_bytes.end()); - - std::vector inv_bits; - inv_bits.reserve(is_indices.size()); - for (int i : is_indices) - inv_bits.push_back(is_inv_flags[i]); - auto inv_mask = pack_bits(inv_bits, 0, (int)inv_bits.size()); - out.insert(out.end(), inv_mask.begin(), inv_mask.end()); + int k_is = compute_optimal_rice_k(is_zz_vals); + payload_writer.write_bits(k_is & 0x07, 3); + rice_encode(payload_writer, is_zz_vals, k_is); } } @@ -501,9 +558,19 @@ std::vector compress_block_adaptive_joint( std::vector sbr_indices; for (int i = 0; i < sbr_end; ++i) if (!active0[i]) sbr_indices.push_back(i); - sbr_write(out, sbr, sbr_indices, can_state_pred, pred_state); + if (!sbr_indices.empty()) + sbr_write_to_stream(payload_writer, sbr, sbr_indices, can_state_pred, pred_state); } + payload_writer.flush(); + auto payload = payload_writer.data(); + uint16_t payload_len = (uint16_t)payload.size(); + header.push_back(payload_len & 0xFF); + header.push_back((payload_len >> 8) & 0xFF); + + std::vector out = header; + out.insert(out.end(), payload.begin(), payload.end()); + pred_state.prev_scale0 = full_scale0; pred_state.prev_scale1 = full_scale1; pred_state.prev_active0.assign(band_count, 0); @@ -522,30 +589,33 @@ std::vector compress_block_adaptive_joint( for (int i = 0; i < sbr_end && i < band_count; ++i) { if (!active0[i]) { pred_state.prev_sbr_valid[i] = 1; - pred_state.prev_sbr_rms[i] = sbr.rms_idx[i]; + pred_state.prev_sbr_rms[i] = sbr.rms_idx[i]; pred_state.prev_sbr_noise[i] = sbr.noise_flag[i] ? 1 : 0; } } - pred_state.ready = true; + pred_state.ready = true; pred_state.last_bi = block_index; return out; } -void save_compressed_buffered(const std::vector>& blocks, - const std::vector& block_modes, - const std::string& path, - uint32_t sr, int num_channels, - float target_kbps, - uint32_t total_frames, - size_t write_buffer_blocks = 8192) +void save_compressed_buffered( + const std::vector>& blocks, + const std::vector& block_modes, + const std::string& path, + uint32_t sr, + int num_channels, + float target_kbps, + uint32_t total_frames, + bool speed_mode, + size_t write_buffer_blocks = 8192) { std::ofstream f(path, std::ios::binary); if (!f) throw std::runtime_error("Cannot create output file"); - + f.write("WHA1", 4); - uint8_t version = 21; + uint8_t version = 24; f.write((char*)&version, 1); f.write((char*)&sr, 4); uint8_t ch = (uint8_t)num_channels; @@ -553,16 +623,17 @@ void save_compressed_buffered(const std::vector>& blocks, uint32_t block_count = (uint32_t)blocks.size(); f.write((char*)&block_count, 4); f.write((char*)&target_kbps, 4); - uint8_t block_format_version = 21; + uint8_t block_format_version = 24; f.write((char*)&block_format_version, 1); - f.write((char*)&total_frames, 4); + uint8_t flags = speed_mode ? 0x01 : 0x00; + f.write((char*)&flags, 1); + int mode_bytes = ((int)block_count + 7) / 8; std::vector mode_packed(mode_bytes, 0); for (uint32_t i = 0; i < block_count; ++i) - if (block_modes[i]) - mode_packed[i / 8] |= (1 << (i % 8)); + if (block_modes[i]) mode_packed[i / 8] |= (1 << (i % 8)); f.write((char*)mode_packed.data(), mode_bytes); std::vector write_buffer; @@ -577,8 +648,8 @@ void save_compressed_buffered(const std::vector>& blocks, for (uint32_t i = 0; i < block_count; ++i) { const auto& blk = blocks[i]; - uint32_t len = (uint32_t)blk.size(); - write_buffer.insert(write_buffer.end(), (char*)&len, (char*)&len + 4); + uint16_t len = (uint16_t)blk.size(); + write_buffer.insert(write_buffer.end(), (char*)&len, (char*)&len + 2); write_buffer.insert(write_buffer.end(), blk.begin(), blk.end()); if (write_buffer.size() >= write_buffer_blocks * 4096) flush_buffer(); @@ -586,79 +657,119 @@ void save_compressed_buffered(const std::vector>& blocks, flush_buffer(); f.close(); - int long_blocks = (int)std::count(block_modes.begin(), block_modes.end(), true); + int long_blocks = (int)std::count(block_modes.begin(), block_modes.end(), true); int short_blocks = (int)block_count - long_blocks; - std::cout << "Saved container v21 to " << path << " (" + std::cout << "Saved container v24 to " << path << " (" << blocks.size() << " blocks, " << long_blocks << " long, " - << short_blocks << " short)" << std::endl; + << short_blocks << " short, speed_mode=" << (speed_mode ? "yes" : "no") << ")" << std::endl; } -std::tuple>, uint32_t, int, float, uint32_t, std::vector> -compress_audio_streaming(const std::string& input_path, - float target_kbps_in, - size_t read_buffer_size_samples = 8192) +std::tuple>, uint32_t, int, float, uint32_t, std::vector, bool> +compress_audio_streaming( + const std::string& input_path, + float target_kbps_in, + size_t read_buffer_size_samples = 8192, + bool speed_mode = false) { uint32_t sr, channels; drwav wav; if (!drwav_init_file(&wav, input_path.c_str(), nullptr)) throw std::runtime_error("Cannot open WAV file"); - sr = wav.sampleRate; + sr = wav.sampleRate; channels = wav.channels; uint64_t total_frames = wav.totalPCMFrameCount; int num_channels = (int)channels; - float target_kbps = (target_kbps_in > 0) ? target_kbps_in : SETTINGS.default_target_kbps; + float target_kbps = (target_kbps_in > 0) ? target_kbps_in : SETTINGS.default_target_kbps; float per_channel_kbps = target_kbps / num_channels; - float overlap_factor = (per_channel_kbps <= 0.0f) ? 0.0f : 1.0f; + float overlap_factor = (per_channel_kbps <= 0.0f) ? 0.0f : 1.0f; PWPT wpt; std::vector> blocks_raw; - std::vector block_modes; + std::vector block_modes; PredContext pred_ctx; - PowCodec pow_enc(num_channels, 0.85f, true, target_kbps / float(num_channels)); + std::unique_ptr pow_enc; + if (!speed_mode) { + pow_enc = std::make_unique(num_channels, 0.85f, true, target_kbps / float(num_channels)); + } - int current_pos = 0; - int prev_block_size = 1024; + int current_pos = 0; + int prev_block_size = 1024; float bits_per_sample = target_kbps * 1000.0f / sr; - int reservoir = 0; - int max_reservoir = (int)(bits_per_sample * 2048 * 0.3); + int reservoir = 0; + int max_reservoir = (int)(bits_per_sample * 2048 * 0.3); std::vector audio_buffer; size_t buffer_start = 0; - bool eof = false; - bool pow_flushed = false; + bool eof = false; + bool pow_flushed = false; + + std::vector analyze_buf, block, read_chunk(read_buffer_size_samples * num_channels); + std::vector processed_buf, tail_buf; + + struct WindowCache { + std::vector buf; + int bs = 0; + int lo = -1; + int ro = -1; + }; + WindowCache win_long, win_short; + + auto get_window = [&](int bs, int lo, int ro) -> const std::vector& { + WindowCache& wc = (bs == 2048) ? win_long : win_short; + if (wc.bs != bs || wc.lo != lo || wc.ro != ro) { + wc.buf.assign(bs, 1.0f); + if (lo > 1) { + for (int i = 0; i < lo; ++i) { + float t = (float)i / lo; + wc.buf[i] = t * t * (3.0f - 2.0f * t); + } + } + if (ro > 1) { + for (int i = 0; i < ro; ++i) { + float t = (float)i / ro; + wc.buf[bs - 1 - i] = t * t * (3.0f - 2.0f * t); + } + } + wc.bs = bs; + wc.lo = lo; + wc.ro = ro; + } + return wc.buf; + }; auto ensure_data_available = [&](int needed_pos) { while ((int)(buffer_start + audio_buffer.size() / num_channels) < needed_pos) { if (eof) { if (!pow_flushed) { pow_flushed = true; - std::vector tail; - pow_enc.flush(tail); - audio_buffer.insert(audio_buffer.end(), tail.begin(), tail.end()); - if (!tail.empty()) continue; + if (!speed_mode && pow_enc) { + tail_buf.clear(); + pow_enc->flush(tail_buf); + audio_buffer.insert(audio_buffer.end(), tail_buf.begin(), tail_buf.end()); + if (!tail_buf.empty()) continue; + } } break; } - size_t samples_to_read = read_buffer_size_samples * num_channels; - std::vector chunk(samples_to_read); - drwav_uint64 frames_read = - drwav_read_pcm_frames_f32(&wav, read_buffer_size_samples, chunk.data()); - + drwav_uint64 frames_read = drwav_read_pcm_frames_f32( + &wav, read_buffer_size_samples, read_chunk.data()); if (frames_read == 0) { eof = true; continue; } - chunk.resize(frames_read * num_channels); - std::vector processed; - processed.reserve(chunk.size()); - pow_enc.process(chunk.data(), (size_t)frames_read, processed); - audio_buffer.insert(audio_buffer.end(), processed.begin(), processed.end()); + if (!speed_mode) { + processed_buf.clear(); + pow_enc->process(read_chunk.data(), (size_t)frames_read, processed_buf); + audio_buffer.insert(audio_buffer.end(), processed_buf.begin(), processed_buf.end()); + } else { + audio_buffer.insert(audio_buffer.end(), read_chunk.begin(), read_chunk.begin() + frames_read * num_channels); + } } }; @@ -667,81 +778,61 @@ compress_audio_streaming(const std::string& input_path, while (current_pos < (int)total_frames) { int analyze_len = std::min(2048, (int)total_frames - current_pos); ensure_data_available(current_pos + analyze_len); - if ((int)(buffer_start + audio_buffer.size() / num_channels) <= current_pos) break; - std::vector analyze_buf(analyze_len * num_channels); + analyze_buf.assign(analyze_len * num_channels, 0.0f); size_t offset = (current_pos - (int)buffer_start) * num_channels; - if (offset + analyze_buf.size() <= audio_buffer.size()) { - std::copy(audio_buffer.begin() + offset, - audio_buffer.begin() + offset + analyze_buf.size(), - analyze_buf.begin()); + std::copy( + audio_buffer.begin() + offset, + audio_buffer.begin() + offset + analyze_buf.size(), + analyze_buf.begin()); } else { size_t avail = audio_buffer.size() - offset; - std::copy(audio_buffer.begin() + offset, - audio_buffer.end(), - analyze_buf.begin()); - std::fill(analyze_buf.begin() + avail / num_channels, - analyze_buf.end(), 0.0f); + std::copy(audio_buffer.begin() + offset, audio_buffer.end(), analyze_buf.begin()); } - bool is_transient = false; - if (num_channels == 2) - is_transient = detect_transient_block_stereo(analyze_buf, num_channels); - else - is_transient = detect_transient_block_mono(analyze_buf, num_channels); + bool is_transient = (num_channels == 2) + ? detect_transient_block_stereo(analyze_buf, num_channels) + : detect_transient_block_mono(analyze_buf, num_channels); int block_size, right_overlap, level; if (is_transient) { - block_size = 1024; + block_size = 1024; right_overlap = (int)(48 * overlap_factor); - level = 4; + level = 4; } else { - block_size = 2048; + block_size = 2048; right_overlap = (int)(96 * overlap_factor); - level = 5; + level = 5; } - int left_overlap = - (prev_block_size == 2048) ? (int)(96 * overlap_factor) - : (int)(48 * overlap_factor); + int left_overlap = (prev_block_size == 2048) + ? (int)(96 * overlap_factor) + : (int)(48 * overlap_factor); int hop = block_size - right_overlap; ensure_data_available(current_pos + block_size); - std::vector block(block_size * num_channels, 0.0f); + block.assign(block_size * num_channels, 0.0f); size_t offset_block = (current_pos - (int)buffer_start) * num_channels; - size_t avail_block = audio_buffer.size() - offset_block; - + size_t avail_block = audio_buffer.size() - offset_block; if (avail_block > 0) { size_t copy_samples = std::min(avail_block, block.size()); - std::copy(audio_buffer.begin() + offset_block, - audio_buffer.begin() + offset_block + copy_samples, - block.begin()); - } - - std::vector window(block_size, 1.0f); - if (left_overlap > 1) { - for (int i = 0; i < left_overlap; ++i) { - float t = (float)i / left_overlap; - window[i] = t * t * (3.0f - 2.0f * t); - } - } - if (right_overlap > 1) { - for (int i = 0; i < right_overlap; ++i) { - float t = (float)i / right_overlap; - window[block_size - 1 - i] = t * t * (3.0f - 2.0f * t); - } + std::copy( + audio_buffer.begin() + offset_block, + audio_buffer.begin() + offset_block + copy_samples, + block.begin()); } - for (int c = 0; c < num_channels; ++c) + const std::vector& window = get_window(block_size, left_overlap, right_overlap); + for (int c = 0; c < num_channels; ++c) { for (int i = 0; i < block_size; ++i) block[i * num_channels + c] *= window[i]; + } bool block_is_full = (current_pos + block_size <= (int)total_frames); - int target_bits = (int)(bits_per_sample * hop); if (target_bits < 256) target_bits = 256; int effective_budget = target_bits + reservoir; @@ -749,36 +840,33 @@ compress_audio_streaming(const std::string& input_path, effective_budget = target_bits + max_reservoir; uint32_t block_index = (uint32_t)blocks_raw.size(); - auto comp = compress_block_adaptive_joint( - block, wpt, num_channels, level, - effective_budget, target_kbps, sr, block_size, - SETTINGS.enable_ms, block_is_full, - block_index, pred_ctx); + block, wpt, num_channels, level, effective_budget, target_kbps, sr, + block_size, SETTINGS.enable_ms, block_is_full, block_index, pred_ctx); int real_bits = (int)comp.size() * 8; reservoir += target_bits - real_bits; if (reservoir < 0) reservoir = 0; if (reservoir > max_reservoir) reservoir = max_reservoir; - blocks_raw.push_back(comp); + + blocks_raw.push_back(std::move(comp)); block_modes.push_back(!is_transient); if (SETTINGS.verbose && (blocks_raw.size() % 10 == 0)) { std::cout << "[Block " << blocks_raw.size() << "] size=" << block_size << ", mode=" << (is_transient ? "short" : "long") - << ", bytes=" << comp.size() + << ", bytes=" << blocks_raw.back().size() << ", reservoir=" << reservoir << std::endl; } - current_pos += hop; - prev_block_size = block_size; + current_pos += hop; + prev_block_size = block_size; int keep_start = current_pos - left_overlap; if (keep_start < 0) keep_start = 0; size_t remove_samples = (keep_start - (int)buffer_start) * num_channels; if (remove_samples > 0 && remove_samples <= audio_buffer.size()) { - audio_buffer.erase(audio_buffer.begin(), - audio_buffer.begin() + remove_samples); + audio_buffer.erase(audio_buffer.begin(), audio_buffer.begin() + remove_samples); buffer_start = keep_start; } @@ -786,132 +874,131 @@ compress_audio_streaming(const std::string& input_path, } drwav_uninit(&wav); - return {blocks_raw, sr, num_channels, target_kbps, (uint32_t)total_frames, block_modes}; + + return { + std::move(blocks_raw), + sr, + num_channels, + target_kbps, + (uint32_t)total_frames, + std::move(block_modes), + speed_mode + }; } #include -int main(int argc, char** argv) { +int main(int argc, char** argv) +{ if (argc < 2) { - std::cerr << "Usage: ./encoder [bitrate_kbps] " - "[read_buffer_samples] [is_max_base]\n" - " is_max_base: -1 auto, 0 disable adaptive IS, >0 fixed max base\n"; + std::cerr << "Usage: ./encoder [bitrate_kbps] [read_buffer_samples] [is_max_base] [sbr_w0] [sbr_w1] [--speed]\n"; return 1; } - std::string input_file = argv[1]; + std::vector args; + bool speed_mode = false; + for (int i = 1; i < argc; ++i) { + if (std::string(argv[i]) == "--speed") { + speed_mode = true; + } else { + args.push_back(argv[i]); + } + } + + int new_argc = args.size() + 1; + char** new_argv = args.data(); + + std::string input_file = new_argv[0]; std::string out_container = input_file; - for (int64_t i = 0; i < 3; i++) - out_container.pop_back(); + for (int64_t i = 0; i < 3; i++) out_container.pop_back(); out_container += "wha"; float target_kbps = SETTINGS.default_target_kbps; - if (argc >= 3) - target_kbps = std::stof(argv[2]); + if (new_argc >= 3) target_kbps = std::stof(new_argv[1]); size_t read_buffer = 8192; - if (argc >= 4) { + if (new_argc >= 4) { try { - read_buffer = std::stoul(argv[3]); + read_buffer = std::stoul(new_argv[2]); if (read_buffer < 256) read_buffer = 256; } catch (...) { - std::cerr << "Invalid read buffer size, using default 8192" << std::endl; + std::cerr << "Invalid read buffer size\n"; } } - if (argc >= 5) { + if (new_argc >= 5) { try { - int mb = std::stoi(argv[4]); + int mb = std::stoi(new_argv[3]); if (mb == 0) { - SETTINGS.enable_adaptive_is = false; + SETTINGS.enable_adaptive_is = false; SETTINGS.adaptive_is_base_override = -1; } else if (mb > 0) { - SETTINGS.enable_adaptive_is = true; + SETTINGS.enable_adaptive_is = true; SETTINGS.adaptive_is_base_override = mb; } else { - SETTINGS.enable_adaptive_is = true; + SETTINGS.enable_adaptive_is = true; SETTINGS.adaptive_is_base_override = -1; } } catch (...) { - std::cerr << "Invalid is_max_base, using auto" << std::endl; - SETTINGS.enable_adaptive_is = true; + SETTINGS.enable_adaptive_is = true; SETTINGS.adaptive_is_base_override = -1; } } + if (new_argc >= 6) { try { SBR_W0 = std::stof(new_argv[4]); } catch (...) {} } + if (new_argc >= 7) { try { SBR_W1 = std::stof(new_argv[5]); } catch (...) {} } + + std::cout << "SBR Shape Weights: W0=" << SBR_W0 + << ", W1=" << SBR_W1 + << " (Contrast: " << (20.0f * std::log10(SBR_W1 / SBR_W0)) << " dB)\n"; + + if (speed_mode) { + std::cout << ">>> SPEED MODE ENABLED: Pre-processing DISABLED (Flag will be saved to container) <<<\n"; + } + struct rusage usage_before, usage_after; getrusage(RUSAGE_SELF, &usage_before); - auto comp = compress_audio_streaming(input_file, target_kbps, read_buffer); - auto blocks = std::get<0>(comp); - uint32_t sr = std::get<1>(comp); - int num_channels = std::get<2>(comp); - float tk = std::get<3>(comp); + auto comp = compress_audio_streaming(input_file, target_kbps, read_buffer, speed_mode); + auto blocks = std::get<0>(comp); + uint32_t sr = std::get<1>(comp); + int num_channels = std::get<2>(comp); + float tk = std::get<3>(comp); uint32_t total_samples = std::get<4>(comp); - auto block_modes = std::get<5>(comp); - - int block_count = (int)blocks.size(); - save_compressed_buffered(blocks, block_modes, out_container, - sr, num_channels, tk, read_buffer); - - getrusage(RUSAGE_SELF, &usage_after); - double user_time_sec = - (usage_after.ru_utime.tv_sec - usage_before.ru_utime.tv_sec) + - (usage_after.ru_utime.tv_usec - usage_before.ru_utime.tv_usec) / 1000000.0; - - double audio_duration_sec = - (total_samples > 0) ? (double)total_samples / (double)sr : 0.0; - - size_t total_file_bytes = 22; - for (const auto& blk : blocks) - total_file_bytes += 4 + blk.size(); - - double actual_bitrate_kbps = - (audio_duration_sec > 0.0) - ? (double)total_file_bytes * 8.0 / 1000.0 / audio_duration_sec - : 0.0; - - double realtime_speed = - (user_time_sec > 0.0) ? (audio_duration_sec / user_time_sec) : 0.0; - - std::cout << "Compressed: " << block_count - << " blocks, user time=" << user_time_sec << "s" << std::endl; - std::cout << "Audio duration: " << audio_duration_sec << " s" << std::endl; - std::cout << "File size: " << total_file_bytes << " bytes" << std::endl; - std::cout << "Target bitrate: " << target_kbps << " kbps" << std::endl; - std::cout << "Actual bitrate: " << actual_bitrate_kbps << " kbps" << std::endl; - std::cout << "Encoding speed: " << realtime_speed << "x realtime" << std::endl; + auto block_modes = std::get<5>(comp); + bool final_speed_mode = std::get<6>(comp); - save_compressed_buffered(blocks, block_modes, out_container, - sr, num_channels, tk, total_samples, read_buffer); + save_compressed_buffered( + blocks, block_modes, out_container, sr, num_channels, tk, + total_samples, final_speed_mode, read_buffer); getrusage(RUSAGE_SELF, &usage_after); double user_time_sec = - (usage_after.ru_utime.tv_sec - usage_before.ru_utime.tv_sec) + + (usage_after.ru_utime.tv_sec - usage_before.ru_utime.tv_sec) + (usage_after.ru_utime.tv_usec - usage_before.ru_utime.tv_usec) / 1000000.0; - double audio_duration_sec = - (total_samples > 0) ? (double)total_samples / (double)sr : 0.0; + double audio_duration_sec = (total_samples > 0) + ? (double)total_samples / (double)sr + : 0.0; - size_t total_file_bytes = 22; + size_t total_file_bytes = 23; for (const auto& blk : blocks) - total_file_bytes += 4 + blk.size(); - - double actual_bitrate_kbps = - (audio_duration_sec > 0.0) - ? (double)total_file_bytes * 8.0 / 1000.0 / audio_duration_sec - : 0.0; - - double realtime_speed = - (user_time_sec > 0.0) ? (audio_duration_sec / user_time_sec) : 0.0; - - std::cout << "Compressed: " << block_count - << " blocks, user time=" << user_time_sec << "s" << std::endl; - std::cout << "Audio duration: " << audio_duration_sec << " s" << std::endl; - std::cout << "File size: " << total_file_bytes << " bytes" << std::endl; - std::cout << "Target bitrate: " << target_kbps << " kbps" << std::endl; - std::cout << "Actual bitrate: " << actual_bitrate_kbps << " kbps" << std::endl; - std::cout << "Encoding speed: " << realtime_speed << "x realtime" << std::endl; + total_file_bytes += 2 + blk.size(); + + double actual_bitrate_kbps = (audio_duration_sec > 0.0) + ? (double)total_file_bytes * 8.0 / 1000.0 / audio_duration_sec + : 0.0; + + double realtime_speed = (user_time_sec > 0.0) + ? (audio_duration_sec / user_time_sec) + : 0.0; + + std::cout << "Compressed: " << blocks.size() << " blocks, user time=" << user_time_sec << "s\n"; + std::cout << "Audio duration: " << audio_duration_sec << " s\n"; + std::cout << "File size: " << total_file_bytes << " bytes\n"; + std::cout << "Target bitrate: " << target_kbps << " kbps\n"; + std::cout << "Actual bitrate: " << actual_bitrate_kbps << " kbps\n"; + std::cout << "Encoding speed: " << realtime_speed << "x realtime\n"; return 0; } diff --git a/coder/joint_stereo.h b/coder/joint_stereo.h index a4b43a5..c8527d6 100644 --- a/coder/joint_stereo.h +++ b/coder/joint_stereo.h @@ -1,6 +1,5 @@ #ifndef JOINT_STEREO_H #define JOINT_STEREO_H - #include #include #include @@ -54,30 +53,21 @@ inline float compute_r_segment(const std::vector& left, const std::vector int start, int end) { double eL = 0.0, eR = 0.0; - for (int i = start; i < end; ++i) { eL += left[i] * left[i]; eR += right[i] * right[i]; } - const double eps = 1e-12; double total = eL + eR; - - if (total < eps) - return 0.5f; - + if (total < eps) return 0.5f; return static_cast(eR / total); } inline bool should_use_adaptive_panorama(const std::array& r_segments) { int changes = 0; - for (int i = 1; i < 4; ++i) { - if (std::abs(r_segments[i] - r_segments[i - 1]) > 0.3f) { - changes++; - } + if (std::abs(r_segments[i] - r_segments[i - 1]) > 0.3f) changes++; } - return (changes >= 2); } @@ -85,91 +75,78 @@ inline void compute_is_parameters_ex(const std::vector& left, const std::vector& right, std::vector& Y, std::array& r_out, - bool& inv_flag, + std::array& inv_flags, bool& use_segmented, int band_idx, int total_bands) { const size_t n = left.size(); - double eL_total = 0.0, eR_total = 0.0, dot_total = 0.0; - for (size_t i = 0; i < n; ++i) { eL_total += left[i] * left[i]; eR_total += right[i] * right[i]; dot_total += left[i] * right[i]; } + + for (int s = 0; s < 4; ++s) inv_flags[s] = false; const double eps = 1e-12; double total = eL_total + eR_total; - if (total < eps || eL_total < eps || eR_total < eps) { - for (int s = 0; s < 4; ++s) - r_out[s] = 0.5f; - + for (int s = 0; s < 4; ++s) r_out[s] = 0.5f; Y.assign(n, 0.0f); - inv_flag = false; use_segmented = false; return; } - double corr = dot_total / (std::sqrt(eL_total * eR_total) + eps); - inv_flag = (corr < -0.2); - - int segmented_threshold = 6 * total_bands / 16; + int segmented_threshold = 8 * total_bands / 16; use_segmented = (band_idx < segmented_threshold); if (use_segmented) { const int seg_size = (int)(n / 4); Y.resize(n); - for (int s = 0; s < 4; ++s) { int start = s * seg_size; int end = (s == 3) ? (int)n : start + seg_size; int len = end - start; - - double seg_eL = 0.0, seg_eR = 0.0; - + + double seg_eL = 0.0, seg_eR = 0.0, seg_dot = 0.0; for (int i = start; i < end; ++i) { seg_eL += left[i] * left[i]; seg_eR += right[i] * right[i]; + seg_dot += left[i] * right[i]; } - + double seg_total = seg_eL + seg_eR; - if (seg_total < eps) { r_out[s] = 0.5f; - for (int i = start; i < end; ++i) - Y[i] = 0.0f; + for (int i = start; i < end; ++i) Y[i] = 0.0f; continue; } - + + double seg_corr = seg_dot / (std::sqrt(seg_eL * seg_eR) + eps); + inv_flags[s] = (seg_corr < -0.2); + r_out[s] = static_cast(seg_eR / seg_total); - double sqrt_eL = std::sqrt(seg_eL); double sqrt_eR = std::sqrt(seg_eR); double inv_sqrt_total = 1.0 / std::sqrt(seg_total); - + double eX = 0.0; std::vector X_seg(len); - for (int i = 0; i < len; ++i) { int idx = start + i; - - double Rterm = inv_flag ? -sqrt_eR * right[idx] : sqrt_eR * right[idx]; + double Rterm = inv_flags[s] ? -sqrt_eR * right[idx] : sqrt_eR * right[idx]; double x = (sqrt_eL * left[idx] + Rterm) * inv_sqrt_total; - X_seg[i] = static_cast(x); eX += x * x; } - + if (eX < eps) { - for (int i = start; i < end; ++i) - Y[i] = 0.0f; + for (int i = start; i < end; ++i) Y[i] = 0.0f; } else { double inv_sqrt_eX = 1.0 / std::sqrt(eX); double target_energy = std::sqrt(seg_total); - for (int i = 0; i < len; ++i) { Y[start + i] = static_cast(X_seg[i] * inv_sqrt_eX * target_energy); } @@ -177,33 +154,31 @@ inline void compute_is_parameters_ex(const std::vector& left, } } else { use_segmented = false; - + double corr = dot_total / (std::sqrt(eL_total * eR_total) + eps); + inv_flags[0] = (corr < -0.2); + float r = static_cast(eR_total / total); - for (int s = 0; s < 4; ++s) - r_out[s] = r; - + for (int s = 0; s < 4; ++s) r_out[s] = r; + double sqrtE = std::sqrt(total); double sqrt_eL = std::sqrt(eL_total); double sqrt_eR = std::sqrt(eR_total); double inv_sqrt_total = 1.0 / std::sqrt(total); - + std::vector X(n); double eX = 0.0; - for (size_t i = 0; i < n; ++i) { - double Rterm = inv_flag ? -sqrt_eR * right[i] : sqrt_eR * right[i]; + double Rterm = inv_flags[0] ? -sqrt_eR * right[i] : sqrt_eR * right[i]; double x = (sqrt_eL * left[i] + Rterm) * inv_sqrt_total; - X[i] = static_cast(x); eX += x * x; } - + if (eX < eps) { Y.assign(n, 0.0f); } else { double inv_sqrt_eX = 1.0 / std::sqrt(eX); Y.resize(n); - for (size_t i = 0; i < n; ++i) { Y[i] = static_cast(X[i] * inv_sqrt_eX * sqrtE); } @@ -213,7 +188,7 @@ inline void compute_is_parameters_ex(const std::vector& left, inline void apply_is(const std::vector& Y, const std::array& r_vals, - bool inv_flag, + const std::array& inv_flags, bool use_segmented, std::vector& left, std::vector& right) @@ -221,31 +196,30 @@ inline void apply_is(const std::vector& Y, size_t n = Y.size(); left.resize(n); right.resize(n); - + if (use_segmented && n >= 16) { int seg_size = (int)(n / 4); - for (int s = 0; s < 4; ++s) { float r = r_vals[s]; float gl = std::sqrt(1.0f - r); float gr = std::sqrt(r); - int start = s * seg_size; int end = (s == 3) ? (int)n : start + seg_size; - + + bool local_inv = inv_flags[s]; for (int i = start; i < end; ++i) { left[i] = Y[i] * gl; - right[i] = Y[i] * (inv_flag ? -gr : gr); + right[i] = Y[i] * (local_inv ? -gr : gr); } } } else { float r = r_vals[0]; float gl = std::sqrt(1.0f - r); float gr = std::sqrt(r); - + bool local_inv = inv_flags[0]; for (size_t i = 0; i < n; ++i) { left[i] = Y[i] * gl; - right[i] = Y[i] * (inv_flag ? -gr : gr); + right[i] = Y[i] * (local_inv ? -gr : gr); } } } @@ -253,20 +227,16 @@ inline void apply_is(const std::vector& Y, inline uint32_t quantize_r(float r, int bits) { float r_clamped = std::clamp(r, 0.0f, 1.0f); float theta = std::asin(std::sqrt(r_clamped)); - int max_val = (1 << bits) - 2; float theta_norm = theta / (float)(M_PI * 0.5f); - return static_cast(std::clamp(theta_norm, 0.0f, 1.0f) * max_val + 0.5f); } inline float dequantize_r(uint32_t q, int bits) { int max_val = (1 << bits) - 2; - float theta_norm = (float)q / max_val; float theta = theta_norm * (float)(M_PI * 0.5f); float s = std::sin(theta); - return s * s; } @@ -288,37 +258,22 @@ inline std::pair, std::vector> mid_side(const std::vec inline bool use_mid_side(float El, float Er, float Em, float Es, bool enable_ms, float target_kbps) { - if (!enable_ms) - return false; - + if (!enable_ms) return false; const float eps = 1e-12f; - - if (target_kbps < 128.0f) { - if (Es > Em * 0.9f) - return false; - + if (target_kbps < 64.0f) { + if (Es > Em * 0.9f) return false; float prod_lr = El * Er; float prod_ms = Em * Es; - - if (prod_ms > prod_lr * 1.05f) - return false; - + if (prod_ms > prod_lr * 1.05f) return false; return true; } - float min_e = std::min(El, Er); float max_e = std::max(El, Er); float energy_ratio = max_e / (min_e + eps); - - if (energy_ratio > 4.0f) - return false; - - if (Es > Em * 0.5f) - return false; - + if (energy_ratio > 4.0f) return false; + if (Es > Em * 0.5f) return false; float prod_lr = El * Er; float prod_ms = Em * Es; - return prod_ms < prod_lr * 0.95f; } @@ -331,81 +286,56 @@ inline bool should_use_is_band(const std::vector& left, float equal_ratio_threshold_low_bitrate) { const size_t n = std::min(left.size(), right.size()); - - if (n == 0) - return false; - + if (n == 0) return false; const double eps = 1e-12; - - double eL = 0.0; - double eR = 0.0; - double cross = 0.0; - + double eL = 0.0, eR = 0.0, cross = 0.0; for (size_t i = 0; i < n; ++i) { const double L = (double)left[i]; const double R = (double)right[i]; - - eL += L * L; - eR += R * R; - cross += L * R; + eL += L * L; eR += R * R; cross += L * R; } - - if (eL + eR < eps) - return false; - - const double correlation = - cross / (std::sqrt(eL * eR) + eps); - - const double abs_correlation = - std::abs(correlation); - + if (eL + eR < eps) return false; + const double correlation = cross / (std::sqrt(eL * eR) + eps); + const double abs_correlation = std::abs(correlation); constexpr double IS_CORRELATION_THRESHOLD = 0.40; - - if (abs_correlation >= IS_CORRELATION_THRESHOLD) - return true; - + if (abs_correlation >= IS_CORRELATION_THRESHOLD) return true; + const double meanL = eL / (double)n; const double meanR = eR / (double)n; - const double minE = std::min(meanL, meanR); const double maxE = std::max(meanL, meanR); - - const double ratio = - maxE / (minE + eps); - - if (ratio > (double)strong_ratio_threshold) - return true; - - double eM = 0.0; - double eS = 0.0; - + const double ratio = maxE / (minE + eps); + if (ratio > (double)strong_ratio_threshold) return true; + + double eM = 0.0, eS = 0.0; for (size_t i = 0; i < n; ++i) { - const double L = (double)left[i]; const double R = (double)right[i]; - - const double M = - (L + R) * 0.5; - - const double S = - (L - R) * 0.5; - - eM += M * M; - eS += S * S; + const double M = (L + R) * 0.5; + const double S = (L - R) * 0.5; + eM += M * M; eS += S * S; } + const bool ms_benefit = use_mid_side((float)eL, (float)eR, (float)eM, (float)eS, enable_ms, target_kbps); + return ms_benefit || maxE >= minE * 4.0; +} - const bool ms_benefit = - use_mid_side( - (float)eL, - (float)eR, - (float)eM, - (float)eS, - enable_ms, - target_kbps - ); +inline int32_t quantize_r_centered(float r, int bits) { + float r_clamped = std::clamp(r, 0.0f, 1.0f); + float theta = std::asin(std::sqrt(r_clamped)); + float theta_center = (float)(M_PI * 0.25); + int steps = (1 << (bits - 1)) - 1; + float step_size = theta_center / steps; + int32_t q = (int32_t)std::round((theta - theta_center) / step_size); + return std::clamp(q, -steps, steps); +} - return ms_benefit || - maxE >= minE * 4.0; +inline float dequantize_r_centered(int32_t q, int bits) { + int steps = (1 << (bits - 1)) - 1; + float theta_center = (float)(M_PI * 0.25); + float step_size = theta_center / steps; + float theta = theta_center + q * step_size; + float s = std::sin(theta); + return s * s; } #endif // JOINT_STEREO_H \ No newline at end of file diff --git a/coder/perc_prior.h b/coder/perc_prior.h index e68eafc..8b1eccd 100644 --- a/coder/perc_prior.h +++ b/coder/perc_prior.h @@ -195,24 +195,6 @@ inline std::vector compute_channel_priority( ); } - float cur_peak = max_v; - - float prev_peak = 0.0f; - for (float v : prev) - prev_peak = std::max(prev_peak, std::fabs(v)); - - constexpr float PEAK_THRESHOLD = 0.01f; - - float transient = 0.0f; - - if (cur_peak > PEAK_THRESHOLD && prev_peak > PEAK_THRESHOLD) { - float peak_ratio = cur_peak / prev_peak; - - transient = std::log2(peak_ratio + 1.0f); - - transient = std::clamp(transient, 0.0f, 4.0f); - } - float p_val = 0.35f * predict + 0.55f * sparsity + diff --git a/coder/pred_state.h b/coder/pred_state.h index 15a065f..5492e3b 100644 --- a/coder/pred_state.h +++ b/coder/pred_state.h @@ -8,6 +8,9 @@ struct ModeState { bool ready = false; uint32_t last_bi = 0; + std::vector prev_k_sub0; + std::vector prev_k_sub1; + std::vector prev_scale0; std::vector prev_scale1; std::vector prev_active0; diff --git a/coder/sbr_encode.h b/coder/sbr_encode.h index 8f5cbe0..716ba0d 100644 --- a/coder/sbr_encode.h +++ b/coder/sbr_encode.h @@ -7,30 +7,78 @@ #include "pred_state.h" #include "entropy_encoder.h" +extern float SBR_W0; +extern float SBR_W1; + +inline void get_sbr_shape_weights(uint8_t mask, float w0, float w1, float out[4]) { + float sum = 0.0f; + float v[4]; + for (int i = 0; i < 4; ++i) { + v[i] = (mask & (1 << i)) ? w1 : w0; + sum += v[i]; + } + for (int i = 0; i < 4; ++i) { + out[i] = v[i] / sum; + } +} + struct SBREncodeResult { std::vector noise_flag; std::vector rms_idx; + std::vector shape_mask; int sbr_end = 0; }; inline SBREncodeResult sbr_analyze( const std::vector>& ch0_original, const std::vector& active0, - int sbr_end, int band_count) + int sbr_end, int band_count, + float w0, float w1) { SBREncodeResult result; result.noise_flag.assign(band_count, false); result.rms_idx.assign(band_count, 0); + result.shape_mask.assign(band_count, 0); result.sbr_end = sbr_end; const float eps = 1e-12f; - for (int i = 0; i < sbr_end; ++i) { if (active0[i]) continue; - const auto& orig = ch0_original[i]; if (orig.empty()) continue; - + + int N = orig.size(); + int seg_size = N / 4; + + float seg_energy[4] = {0, 0, 0, 0}; + for (int s = 0; s < 4; ++s) { + int start = s * seg_size; + int end = (s == 3) ? N : (s + 1) * seg_size; + for (int j = start; j < end; ++j) { + seg_energy[s] += orig[j] * orig[j]; + } + } + float total_e = seg_energy[0] + seg_energy[1] + seg_energy[2] + seg_energy[3] + eps; + float target_p[4]; + for(int s = 0; s < 4; ++s) target_p[s] = seg_energy[s] / total_e; + + float best_err = 1e9f; + uint8_t best_mask = 0; + for (uint8_t m = 0; m < 16; ++m) { + float w[4]; + get_sbr_shape_weights(m, w0, w1, w); + float err = 0.0f; + for(int s = 0; s < 4; ++s) { + float diff = target_p[s] - w[s]; + err += diff * diff; + } + if (err < best_err) { + best_err = err; + best_mask = m; + } + } + result.shape_mask[i] = best_mask; + double sum_sq = 0.0; float max_abs = 0.0f; for (float v : orig) { @@ -38,13 +86,9 @@ inline SBREncodeResult sbr_analyze( float av = std::abs(v); if (av > max_abs) max_abs = av; } - float rms = std::sqrt(static_cast(sum_sq / orig.size()) + eps); + float rms = std::sqrt(static_cast(sum_sq / N) + eps); float crest = max_abs / (rms + eps); - - // if (i >= band_count / 2) { - // std::cout << crest << '\n'; - // } - + result.noise_flag[i] = (crest < 3.5f); int sb = get_scale_bits(i); result.rms_idx[i] = get_scale_idx(rms, sb); @@ -52,79 +96,40 @@ inline SBREncodeResult sbr_analyze( return result; } -inline void sbr_write(std::vector& out, - const SBREncodeResult& sbr, - const std::vector& sbr_indices, - bool can_state_pred, - const ModeState& pred_state) +inline void sbr_write_to_stream(BitWriterMSB& writer, + const SBREncodeResult& sbr, + const std::vector& sbr_indices, + bool can_state_pred, + const ModeState& pred_state) { if (sbr_indices.empty()) return; - std::vector noise_bits; - noise_bits.reserve(sbr_indices.size()); for (int idx : sbr_indices) { - bool cur = sbr.noise_flag[idx]; + bool cur = sbr.noise_flag[idx]; bool write_bit = cur; - if (can_state_pred && - idx < (int)pred_state.prev_sbr_valid.size() && - pred_state.prev_sbr_valid[idx]) - { + if (can_state_pred && idx < (int)pred_state.prev_sbr_valid.size() && pred_state.prev_sbr_valid[idx]) { write_bit = cur ^ (pred_state.prev_sbr_noise[idx] != 0); } - noise_bits.push_back(write_bit); + writer.write_bits(write_bit ? 1 : 0, 1); } - auto noise_packed = pack_bits(noise_bits, 0, (int)noise_bits.size()); - out.insert(out.end(), noise_packed.begin(), noise_packed.end()); - - uint8_t sbr_rms_mode = 0; - std::vector sbr_vals; - int k_sbr = 0; - - if (can_state_pred) { - sbr_vals.reserve(sbr_indices.size()); - for (int idx : sbr_indices) { - uint32_t cur = sbr.rms_idx[idx]; - uint32_t val = cur; - if (idx < (int)pred_state.prev_sbr_valid.size() && - pred_state.prev_sbr_valid[idx]) - { - int32_t d = (int32_t)cur - (int32_t)pred_state.prev_sbr_rms[idx]; - val = zigzag_encode(d); - } - sbr_vals.push_back(val); - } - k_sbr = compute_optimal_rice_k(sbr_vals, 7); - - size_t fixed_bits = 0; - for (int idx : sbr_indices) - fixed_bits += get_scale_bits(idx); - size_t rice_bits = estimate_rice_bits(sbr_vals, k_sbr) + 24; - if (rice_bits < fixed_bits) - sbr_rms_mode = 1; + for (int idx : sbr_indices) { + writer.write_bits(sbr.shape_mask[idx], 4); } - out.push_back(sbr_rms_mode); - - if (sbr_rms_mode == 1) { - out.push_back((uint8_t)k_sbr); - BitWriterMSB rms_writer; - for (uint32_t v : sbr_vals) - rice_encode(rms_writer, std::vector{v}, k_sbr); - rms_writer.flush(); - auto rms_bytes = rms_writer.data(); - uint16_t rms_len = (uint16_t)rms_bytes.size(); - out.push_back(rms_len & 0xFF); - out.push_back((rms_len >> 8) & 0xFF); - out.insert(out.end(), rms_bytes.begin(), rms_bytes.end()); - } else { - BitWriterMSB rms_writer; - for (int idx : sbr_indices) { - int sb = get_scale_bits(idx); - rms_writer.write_bits(sbr.rms_idx[idx], sb); + std::vector sbr_vals; + sbr_vals.reserve(sbr_indices.size()); + for (int idx : sbr_indices) { + uint32_t cur = sbr.rms_idx[idx]; + uint32_t val = cur; + if (can_state_pred && idx < (int)pred_state.prev_sbr_valid.size() && pred_state.prev_sbr_valid[idx]) { + int32_t d = (int32_t)cur - (int32_t)pred_state.prev_sbr_rms[idx]; + val = zigzag_encode(d); } - rms_writer.flush(); - auto rms_bytes = rms_writer.data(); - out.insert(out.end(), rms_bytes.begin(), rms_bytes.end()); + sbr_vals.push_back(val); } + + int k_sbr = compute_optimal_rice_k(sbr_vals, 7); + writer.write_bits(k_sbr & 0x07, 3); + rice_encode(writer, sbr_vals, k_sbr); } \ No newline at end of file diff --git a/decoder/codec_utils.h b/decoder/codec_utils.h index 4bd5cd9..55f48f1 100644 --- a/decoder/codec_utils.h +++ b/decoder/codec_utils.h @@ -8,6 +8,7 @@ static constexpr float SCALE_LOG_MIN = -6.0f; static constexpr float SCALE_LOG_MAX = 1.0f; inline int get_scale_bits(int band_idx) { + return 8; if (band_idx < 1) return 12; return 8; } diff --git a/decoder/decoder.cpp b/decoder/decoder.cpp index b56dbba..86577a1 100644 --- a/decoder/decoder.cpp +++ b/decoder/decoder.cpp @@ -8,7 +8,7 @@ #include #include #include -#include +#include #include #define DR_WAV_IMPLEMENTATION #include "dr_wav.h" @@ -23,85 +23,122 @@ #include "sbr_decode.h" #include "overlap_add.h" #include "pow_filter.h" + #ifndef M_PI #define M_PI 3.14159265358979323846 #endif -void decompress_wha_to_wav(const std::string& in_wha, - const std::string& out_wav, - size_t write_batch_frames = 8192) +float SBR_W0 = 1.0f; +float SBR_W1 = 9.0f; + +void decompress_wha_to_wav( + const std::string& in_wha, + const std::string& out_wav, + size_t write_batch_frames = 8192) { std::ifstream f(in_wha, std::ios::binary); if (!f) throw std::runtime_error("Cannot open " + in_wha); - auto magic = read_n(f, 4); - if (magic.size() != 4 || std::string((char*)magic.data(), 4) != "WHA1") - throw std::runtime_error("Not a WHA container"); + auto write_samples = [&](const std::vector& samples) { + if (samples.empty()) return; + size_t frames = samples.size() / num_channels; + if (frames == 0) return; uint8_t version = read_u8(f); - if (version != 18 && version != 19 && version != 20 && version != 21) - throw std::runtime_error("Unsupported container version (only v18/v19/v20/v21)"); - - uint32_t sr = read_u32(f); - uint8_t num_channels = read_u8(f); - uint32_t block_count = read_u32(f); - float target_kbps = read_f32(f); - uint8_t block_format_version = read_u8(f); - - if (block_format_version != 18 && block_format_version != 19 && block_format_version != 20 && block_format_version != 21) - throw std::runtime_error("Unsupported block format version (only v18/v19/v20/v21)"); + if (version < 18 || version > 24) + throw std::runtime_error("Unsupported container version (only v18-v24)"); + + uint32_t sr = read_u32(f); + uint8_t num_channels = read_u8(f); + uint32_t block_count = read_u32(f); + float target_kbps = read_f32(f); + uint8_t block_format_version = read_u8(f); + + if (block_format_version < 18 || block_format_version > 24) + throw std::runtime_error("Unsupported block format version (only v18-v24)"); uint32_t total_frames = 0; - if (version >= 21) { - total_frames = read_u32(f); + if (version >= 21) total_frames = read_u32(f); + + uint8_t flags = 0; + if (version >= 24) { + flags = read_u8(f); + } + + bool speed_mode = (flags & 0x01) != 0; + if (speed_mode) { + std::cout << ">>> Container indicates SPEED MODE: Post-processing will be skipped automatically <<<\n"; } bool pred_enabled = (block_format_version >= 19); + float per_channel_kbps = target_kbps / num_channels; - float overlap_factor = (per_channel_kbps <= 0.0f) ? 0.0f : 1.0f; + float overlap_factor = (per_channel_kbps <= 0.0f) ? 0.0f : 1.0f; int mode_bytes = (block_count + 7) / 8; std::vector mode_packed = read_n(f, mode_bytes); - std::vector block_modes = unpack_bits(mode_packed.data(), mode_bytes, block_count); + std::vector block_modes = unpack_bits( + mode_packed.data(), mode_bytes, block_count); drwav_data_format fmt; - fmt.container = drwav_container_riff; - fmt.format = DR_WAVE_FORMAT_IEEE_FLOAT; - fmt.channels = num_channels; - fmt.sampleRate = sr; + fmt.container = drwav_container_riff; + fmt.format = DR_WAVE_FORMAT_IEEE_FLOAT; + fmt.channels = num_channels; + fmt.sampleRate = sr; fmt.bitsPerSample = 32; drwav wav; if (!drwav_init_file_write(&wav, out_wav.c_str(), &fmt, nullptr)) throw std::runtime_error("dr_wav failed to init"); - PWPT wpt; - ModeState long_state; - ModeState short_state; - bool stereo = (num_channels == 2); + PWPT wpt; + ModeState long_state, short_state; + bool stereo = (num_channels == 2); + + std::unique_ptr pow_dec; + if (!speed_mode) { + pow_dec = std::make_unique((int)num_channels, 0.85f, false, target_kbps / float(num_channels)); + } - PowCodec pow_dec((int)num_channels, 0.85f, false, target_kbps / float(num_channels)); OverlapAddState ola; ola.init(num_channels, write_batch_frames); - std::vector write_buf; + std::vector write_buf, decoded; write_buf.reserve(write_batch_frames * num_channels); + decoded.reserve(write_batch_frames * num_channels); std::unique_ptr postproc; - if (shouldApplySpectralFloor((int)num_channels, target_kbps)) { + if (!speed_mode && shouldApplySpectralFloor((int)num_channels, target_kbps)) { SpectralFloorPostFX::Config pp_cfg; - postproc = std::make_unique((int)num_channels, pp_cfg); + postproc = std::make_unique( + (int)num_channels, pp_cfg); } - std::vector window, window_sq; - std::vector mode_ms, active0, active1; + std::vector band_shapes_long = compute_band_shapes(2048, 5); + std::vector band_shapes_short = compute_band_shapes(1024, 4); + + struct WindowCache { + std::vector window, window_sq; + int bs = 0; + int lo = -1; + int ro = -1; + }; + WindowCache win_long, win_short; + + std::vector blk, mode_ms, active0, active1; std::vector> is_r; - std::vector is_inv, is_use_segmented; + std::vector> is_inv; + std::vector is_use_segmented; std::vector steps0, steps1; std::vector> quant0, quant1; std::vector> ch0_bands, ch1_bands; - std::vector> recon_chs; + quant0.reserve(32); + quant1.reserve(32); + ch0_bands.reserve(32); + ch1_bands.reserve(32); + + std::vector> recon_chs(num_channels); uint64_t frames_written = 0; auto write_samples = [&](const std::vector& samples) { @@ -133,9 +170,24 @@ void decompress_wha_to_wav(const std::string& in_wha, } drwav_write_pcm_frames(&wav, (uint64_t)frames_to_write, data_to_write); - frames_written += frames_to_write; + frames_written += frames_to_write; + }; + + auto flush_write_buffer = [&]() { + if (write_buf.empty()) return; + + if (!speed_mode && pow_dec) { + decoded.clear(); + pow_dec->process(write_buf.data(), write_buf.size() / num_channels, decoded); + write_samples(decoded); + } else { + write_samples(write_buf); + } + write_buf.clear(); }; + int prev_block_size = 1024; + auto flush_write_buffer = [&]() { if (write_buf.empty()) return; std::vector decoded; @@ -148,161 +200,154 @@ void decompress_wha_to_wav(const std::string& in_wha, int prev_block_size = 1024; for (uint32_t bi = 0; bi < block_count; ++bi) { bool use_long_block = block_modes[bi]; - int level = use_long_block ? 5 : 4; - int block_size = use_long_block ? 2048 : 1024; - int right_overlap = use_long_block ? (int)(96 * overlap_factor) : (int)(48 * overlap_factor); - int left_overlap = (prev_block_size == 2048) ? (int)(96 * overlap_factor) : (int)(48 * overlap_factor); + int level = use_long_block ? 5 : 4; + int block_size = use_long_block ? 2048 : 1024; + int right_overlap = use_long_block + ? (int)(96 * overlap_factor) + : (int)(48 * overlap_factor); + int left_overlap = (prev_block_size == 2048) + ? (int)(96 * overlap_factor) + : (int)(48 * overlap_factor); if (bi == 0) left_overlap = 0; - int hop = block_size - right_overlap; + int hop = block_size - right_overlap; - int total_bands = 1 << level; - std::vector band_shapes = compute_band_shapes(block_size, level); - int expected_band_count = total_bands; + int total_bands = 1 << level; + const std::vector& band_shapes = use_long_block + ? band_shapes_long + : band_shapes_short; + int expected_band_count = total_bands; ModeState& pred_state = use_long_block ? long_state : short_state; pred_state.ensure(expected_band_count); bool can_state_pred = false; if (pred_enabled && pred_state.ready && bi > pred_state.last_bi) { uint32_t age = bi - pred_state.last_bi; - if (age <= SAME_MODE_PRED_MAX_AGE) can_state_pred = true; + if (age <= SAME_MODE_PRED_MAX_AGE) + can_state_pred = true; } - window.clear(); window_sq.clear(); - make_window(window, window_sq, block_size, left_overlap, right_overlap); + WindowCache& wc = use_long_block ? win_long : win_short; + if (wc.bs != block_size || wc.lo != left_overlap || wc.ro != right_overlap) { + make_window(wc.window, wc.window_sq, block_size, left_overlap, right_overlap); + wc.bs = block_size; + wc.lo = left_overlap; + wc.ro = right_overlap; + } - uint32_t block_len = read_u32(f); - std::vector blk(block_len); - f.read((char*)blk.data(), block_len); - if ((size_t)f.gcount() < block_len) throw std::runtime_error("Block truncated"); - - size_t ptr = 0; - auto need = [&](size_t n) { if (ptr + n > blk.size()) throw std::runtime_error("Block truncated"); }; + const std::vector& window = wc.window; + const std::vector& window_sq = wc.window_sq; - int mode_bytes_loc = (expected_band_count + 7) / 8; - mode_ms.resize(expected_band_count); - for (int i = 0; i < expected_band_count; ++i) { - uint8_t byte = blk[ptr + (i >> 3)]; mode_ms[i] = (byte >> (i & 7)) & 1; + size_t block_len; + if (block_format_version >= 23) { + uint16_t len16; + f.read((char*)&len16, 2); + block_len = len16; + } else { + block_len = read_u32(f); } - ptr += mode_bytes_loc; - int mask_bytes = (expected_band_count + 7) >> 3; - active0.resize(expected_band_count); - for (int i = 0; i < expected_band_count; ++i) { - uint8_t byte = blk[ptr + (i >> 3)]; active0[i] = (byte >> (i & 7)) & 1; - } - ptr += mask_bytes; + blk.resize(block_len); + f.read((char*)blk.data(), block_len); + if ((size_t)f.gcount() < block_len) + throw std::runtime_error("Block truncated"); - if (stereo) { - active1.resize(expected_band_count); + size_t ptr = 0; + auto need = [&](size_t n) { + if (ptr + n > blk.size()) + throw std::runtime_error("Block truncated"); + }; + + if (block_format_version >= 24) { + mode_ms.resize(expected_band_count); + active0.resize(expected_band_count); + if (stereo) active1.resize(expected_band_count); + } else { + mode_ms.resize(expected_band_count); + int mode_bytes_loc = (expected_band_count + 7) / 8; for (int i = 0; i < expected_band_count; ++i) { - uint8_t byte = blk[ptr + (i >> 3)]; active1[i] = (byte >> (i & 7)) & 1; + uint8_t byte = blk[ptr + (i >> 3)]; + mode_ms[i] = (byte >> (i & 7)) & 1; + } + ptr += mode_bytes_loc; + + int mask_bytes = (expected_band_count + 7) >> 3; + active0.resize(expected_band_count); + for (int i = 0; i < expected_band_count; ++i) { + uint8_t byte = blk[ptr + (i >> 3)]; + active0[i] = (byte >> (i & 7)) & 1; } ptr += mask_bytes; + + if (stereo) { + active1.resize(expected_band_count); + for (int i = 0; i < expected_band_count; ++i) { + uint8_t byte = blk[ptr + (i >> 3)]; + active1[i] = (byte >> (i & 7)) & 1; + } + ptr += mask_bytes; + } } need(1); uint8_t k_scale_byte = blk[ptr++]; - int k_scale0 = k_scale_byte & 0x07; - int k_scale1 = stereo ? ((k_scale_byte >> 3) & 0x07) : 0; - bool block_is_full = (k_scale_byte & 0x80) != 0; + int k_scale0 = k_scale_byte & 0x07; + int k_scale1 = stereo ? ((k_scale_byte >> 3) & 0x07) : 0; + bool block_is_full = (k_scale_byte & 0x80) != 0; bool use_scale_pred = pred_enabled && can_state_pred && ((k_scale_byte & 0x40) != 0); - int sbr_end = block_is_full ? (3 * expected_band_count / 4) : 0; + int sbr_end = block_is_full ? (3 * expected_band_count / 4) : 0; need(2); uint16_t payload_len; memcpy(&payload_len, &blk[ptr], 2); ptr += 2; - if (ptr + payload_len > blk.size()) throw std::runtime_error("Payload truncated"); + if (ptr + payload_len > blk.size()) + throw std::runtime_error("Payload truncated"); BitReaderMSB payload_reader(blk.data() + ptr, payload_len); ptr += payload_len; - is_r.resize(expected_band_count); - for (auto& a : is_r) a.fill(0.5f); - is_inv.resize(expected_band_count, false); - is_use_segmented.resize(expected_band_count, false); - bool is_used = false; - - if (block_format_version >= 18 && stereo && target_kbps < 510.0f) { - int is_start = get_is_start_band(target_kbps, total_bands); - if (is_start < expected_band_count) { - int segmented_threshold = 6 * total_bands / 16; - if (block_format_version >= 20) { - std::vector is_indices; - for (int i = is_start; i < expected_band_count; ++i) if (mode_ms[i]) is_indices.push_back(i); - int total_bits = 0; - for (int idx : is_indices) { - int bits = get_r_bits(idx); - total_bits += (idx < segmented_threshold) ? bits * 4 : bits; - } - int total_bytes = (total_bits + 7) / 8; - int inv_bytes = ((int)is_indices.size() + 7) / 8; - if (ptr + total_bytes + inv_bytes <= blk.size()) { - if (!is_indices.empty()) { - is_used = true; - BitReaderMSB r_reader(blk.data() + ptr, total_bytes); - for (int idx : is_indices) { - int bits = get_r_bits(idx); - if (idx < segmented_threshold) { - is_use_segmented[idx] = true; - for (int s = 0; s < 4; ++s) { - uint32_t q = r_reader.read_bits(bits); - is_r[idx][s] = dequantize_r(q, bits); - } - } else { - is_use_segmented[idx] = false; - uint32_t q = r_reader.read_bits(bits); - is_r[idx].fill(dequantize_r(q, bits)); - } - } - ptr += total_bytes; - auto inv_tmp = unpack_bits(blk.data() + ptr, inv_bytes, (int)is_indices.size()); - for (size_t j = 0; j < is_indices.size(); ++j) is_inv[is_indices[j]] = inv_tmp[j]; - ptr += inv_bytes; - } - } - } else { - int total_bits = 0; - for (int i = is_start; i < expected_band_count; ++i) { - int bits = get_r_bits(i); - total_bits += (i < segmented_threshold) ? bits * 4 : bits; - } - int total_bytes = (total_bits + 7) / 8; - int inv_bytes = (expected_band_count - is_start + 7) / 8; - if (ptr + total_bytes + inv_bytes <= blk.size()) { - is_used = true; - BitReaderMSB r_reader(blk.data() + ptr, total_bytes); - for (int i = is_start; i < expected_band_count; ++i) { - int bits = get_r_bits(i); - if (i < segmented_threshold) { - is_use_segmented[i] = true; - for (int s = 0; s < 4; ++s) { - uint32_t q = r_reader.read_bits(bits); - is_r[i][s] = dequantize_r(q, bits); - } - } else { - is_use_segmented[i] = false; - uint32_t q = r_reader.read_bits(bits); - is_r[i].fill(dequantize_r(q, bits)); - } - } - ptr += total_bytes; - auto inv_tmp = unpack_bits(blk.data() + ptr, inv_bytes, expected_band_count - is_start); - for (int i = is_start; i < expected_band_count; ++i) is_inv[i] = inv_tmp[i - is_start]; - ptr += inv_bytes; + if (block_format_version >= 24) { + uint32_t packed_mode = payload_reader.read_bits(expected_band_count); + uint32_t packed_active0 = payload_reader.read_bits(expected_band_count); + for (int i = 0; i < expected_band_count; ++i) { + mode_ms[i] = (packed_mode >> (expected_band_count - 1 - i)) & 1; + active0[i] = (packed_active0 >> (expected_band_count - 1 - i)) & 1; + } + if (stereo) { + int is_start = get_is_start_band(target_kbps, total_bands); + int non_is_count = 0; + for (int i = 0; i < expected_band_count; ++i) { + if (!((i >= is_start) && mode_ms[i])) non_is_count++; + } + uint32_t packed_active1 = 0; + if (non_is_count > 0) { + packed_active1 = payload_reader.read_bits(non_is_count); + } + int bit_idx = non_is_count - 1; + for (int i = 0; i < expected_band_count; ++i) { + bool is_is_band = (i >= is_start) && mode_ms[i]; + if (!is_is_band) { + active1[i] = (packed_active1 >> bit_idx) & 1; + bit_idx--; + } else { + active1[i] = 0; } } } } - SBRDecodeResult sbr = sbr_decode(blk.data(), ptr, blk.size(), active0, sbr_end, expected_band_count, block_format_version, can_state_pred, pred_state); - - steps0.resize(expected_band_count, 0.0f); - steps1.resize(expected_band_count, 0.0f); + steps0.assign(expected_band_count, 0.0f); + steps1.assign(expected_band_count, 0.0f); - auto decode_scales = [&](std::vector& steps, const std::vector& active, int k_scale, std::vector& prev_scale, const std::vector& prev_active) { + auto decode_scales = [&]( + std::vector& steps, + const std::vector& active, + int k_scale, + std::vector& prev_scale, + const std::vector& prev_active) + { for (int i = 0; i < expected_band_count; ++i) { if (!active[i]) continue; auto idx_vec = rice_decode(payload_reader, 1, k_scale); @@ -312,6 +357,7 @@ void decompress_wha_to_wav(const std::string& in_wha, int max_idx = (1 << sb) - 1; uint32_t idx; if (use_scale_pred && i < (int)prev_active.size() && prev_active[i]) { + if (i >= (int)prev_scale.size()) throw std::runtime_error("prev_scale out of bounds"); int32_t d = zigzag_decode(val); int64_t v = (int64_t)prev_scale[i] + d; if (v < 0) v = 0; @@ -327,25 +373,208 @@ void decompress_wha_to_wav(const std::string& in_wha, }; decode_scales(steps0, active0, k_scale0, pred_state.prev_scale0, pred_state.prev_active0); - if (stereo) decode_scales(steps1, active1, k_scale1, pred_state.prev_scale1, pred_state.prev_active1); + if (stereo) + decode_scales(steps1, active1, k_scale1, pred_state.prev_scale1, pred_state.prev_active1); + + int k_sub_k0 = (int)payload_reader.read_bits(2); + int k_sub_k1 = stereo ? (int)payload_reader.read_bits(2) : 0; + + int active_count0 = 0, active_count1 = 0; + for (int i = 0; i < expected_band_count; ++i) { + if (active0[i]) active_count0++; + if (stereo && active1[i]) active_count1++; + } + + auto k_sub_vals0 = rice_decode(payload_reader, active_count0, k_sub_k0); + std::vector k_sub_vals1; + if (stereo) + k_sub_vals1 = rice_decode(payload_reader, active_count1, k_sub_k1); quant0.resize(expected_band_count); + quant1.resize(expected_band_count); + ch0_bands.resize(expected_band_count); + ch1_bands.resize(expected_band_count); + + int k_idx0 = 0; for (int i = 0; i < expected_band_count; ++i) { + quant0[i].clear(); if (!active0[i]) continue; - int k_sub = (int)payload_reader.read_bits(3); + if (k_idx0 >= (int)k_sub_vals0.size()) throw std::runtime_error("k_sub_vals0 out of bounds"); + int k_sub = k_sub_vals0[k_idx0++]; + if (k_sub > 7) k_sub = 7; auto uvals = rice_decode(payload_reader, band_shapes[i], k_sub); + if ((int)uvals.size() != band_shapes[i]) throw std::runtime_error("uvals size mismatch ch0"); quant0[i].resize(uvals.size()); - for (size_t j = 0; j < uvals.size(); ++j) quant0[i][j] = zigzag_decode(uvals[j]); + for (size_t j = 0; j < uvals.size(); ++j) + quant0[i][j] = zigzag_decode(uvals[j]); } if (stereo) { - quant1.resize(expected_band_count); + int k_idx1 = 0; for (int i = 0; i < expected_band_count; ++i) { + quant1[i].clear(); if (!active1[i]) continue; - int k_sub = (int)payload_reader.read_bits(3); + if (k_idx1 >= (int)k_sub_vals1.size()) throw std::runtime_error("k_sub_vals1 out of bounds"); + int k_sub = k_sub_vals1[k_idx1++]; + if (k_sub > 7) k_sub = 7; auto uvals = rice_decode(payload_reader, band_shapes[i], k_sub); + if ((int)uvals.size() != band_shapes[i]) throw std::runtime_error("uvals size mismatch ch1"); quant1[i].resize(uvals.size()); - for (size_t j = 0; j < uvals.size(); ++j) quant1[i][j] = zigzag_decode(uvals[j]); + for (size_t j = 0; j < uvals.size(); ++j) + quant1[i][j] = zigzag_decode(uvals[j]); + } + } + + is_r.resize(expected_band_count); + for (auto& a : is_r) a.fill(0.5f); + is_inv.assign(expected_band_count, std::array{}); + for (auto& a : is_inv) a.fill(false); + is_use_segmented.assign(expected_band_count, false); + + bool is_used = false; + if (stereo && target_kbps < 510.0f) { + int is_start = get_is_start_band(target_kbps, total_bands); + if (is_start < expected_band_count) { + std::vector is_indices; + for (int i = is_start; i < expected_band_count; ++i) + if (mode_ms[i]) is_indices.push_back(i); + + if (!is_indices.empty()) { + is_used = true; + int segmented_threshold = 8 * total_bands / 16; + int total_vals = 0; + for (int idx : is_indices) + total_vals += (idx < segmented_threshold) ? 4 : 1; + + if (block_format_version >= 24) { + int total_inv_bits = 0; + for (int idx : is_indices) + total_inv_bits += (idx < segmented_threshold) ? 4 : 1; + + std::vector inv_bits(total_inv_bits); + int abs_idx = 0; + int bits_left = total_inv_bits; + while (bits_left > 0) { + int to_read = std::min(32, bits_left); + uint32_t chunk = payload_reader.read_bits(to_read); + for (int i = 0; i < to_read; ++i) { + inv_bits[abs_idx++] = (chunk >> (to_read - 1 - i)) & 1; + } + bits_left -= to_read; + } + + int read_idx = 0; + for (int idx : is_indices) { + bool is_seg = (idx < segmented_threshold); + if (is_seg) { + for (int s = 0; s < 4; ++s) { + is_inv[idx][s] = (inv_bits[read_idx++] != 0); + } + } else { + is_inv[idx][0] = (inv_bits[read_idx++] != 0); + } + } + } else { + int total_inv_bits = (int)is_indices.size(); + std::vector inv_bits(total_inv_bits); + if (total_inv_bits > 0) { + uint32_t chunk = payload_reader.read_bits(total_inv_bits); + for (int i = 0; i < total_inv_bits; ++i) { + inv_bits[i] = (chunk >> (total_inv_bits - 1 - i)) & 1; + } + } + int read_idx = 0; + for (int idx : is_indices) { + bool inv_val = (inv_bits[read_idx++] != 0); + for (int s = 0; s < 4; ++s) is_inv[idx][s] = inv_val; + } + } + + int k_is = (int)payload_reader.read_bits(3); + auto is_zz_vals = rice_decode(payload_reader, total_vals, k_is); + if ((int)is_zz_vals.size() != total_vals) + throw std::runtime_error("IS Rice decode size mismatch"); + + int val_idx = 0; + for (int idx : is_indices) { + int bits = get_r_bits(idx); + if (idx < segmented_threshold) { + is_use_segmented[idx] = true; + for (int s = 0; s < 4; ++s) { + int32_t q = zigzag_decode(is_zz_vals[val_idx++]); + is_r[idx][s] = dequantize_r_centered(q, bits); + } + } else { + is_use_segmented[idx] = false; + int32_t q = zigzag_decode(is_zz_vals[val_idx++]); + is_r[idx].fill(dequantize_r_centered(q, bits)); + } + } + } + } + } else if (block_format_version >= 18 && block_format_version < 24) { + if (stereo && target_kbps < 510.0f) { + int is_start = get_is_start_band(target_kbps, total_bands); + if (is_start < expected_band_count) { + int segmented_threshold = 8 * total_bands / 16; + if (block_format_version >= 20) { + std::vector is_indices; + for (int i = is_start; i < expected_band_count; ++i) + if (mode_ms[i]) is_indices.push_back(i); + + if (block_format_version >= 22 && !is_indices.empty()) { + is_used = true; + BitReaderMSB is_reader(blk.data() + ptr, blk.size() - ptr); + size_t is_bits_read = 0; + + int total_inv_bits = (int)is_indices.size(); + std::vector inv_bits(total_inv_bits); + if (total_inv_bits > 0) { + uint32_t chunk = is_reader.read_bits(total_inv_bits); + for (int i = 0; i < total_inv_bits; ++i) { + inv_bits[i] = (chunk >> (total_inv_bits - 1 - i)) & 1; + } + } + int read_idx = 0; + for (int idx : is_indices) { + bool inv_val = (inv_bits[read_idx++] != 0); + for (int s = 0; s < 4; ++s) is_inv[idx][s] = inv_val; + is_bits_read += 1; + } + + int k_is = (int)is_reader.read_bits(3); + is_bits_read += 3; + + int total_vals = 0; + for (int idx : is_indices) + total_vals += (idx < segmented_threshold) ? 4 : 1; + + auto is_zz_vals = rice_decode(is_reader, total_vals, k_is); + if ((int)is_zz_vals.size() != total_vals) + throw std::runtime_error("IS Rice decode size mismatch"); + + for (uint32_t v : is_zz_vals) + is_bits_read += (v >> k_is) + 1 + k_is; + + int val_idx = 0; + for (int idx : is_indices) { + int bits = get_r_bits(idx); + if (idx < segmented_threshold) { + is_use_segmented[idx] = true; + for (int s = 0; s < 4; ++s) { + int32_t q = zigzag_decode(is_zz_vals[val_idx++]); + is_r[idx][s] = dequantize_r_centered(q, bits); + } + } else { + is_use_segmented[idx] = false; + int32_t q = zigzag_decode(is_zz_vals[val_idx++]); + is_r[idx].fill(dequantize_r_centered(q, bits)); + } + } + ptr += (is_bits_read + 7) / 8; + } + } + } } }; @@ -361,18 +590,31 @@ void decompress_wha_to_wav(const std::string& in_wha, for (size_t j = 0; j < uvals.size(); ++j) quant0[i][j] = zigzag_decode(uvals[j]); } - ch0_bands.resize(expected_band_count); - ch1_bands.resize(expected_band_count); + SBRDecodeResult sbr; + if (block_format_version >= 23) { + sbr = sbr_decode_v23( + payload_reader, active0, sbr_end, expected_band_count, + can_state_pred, pred_state); + } else { + sbr = sbr_decode( + blk.data(), ptr, blk.size(), active0, sbr_end, + expected_band_count, block_format_version, + can_state_pred, pred_state); + } + for (int i = 0; i < expected_band_count; ++i) { + ch0_bands[i].clear(); + ch0_bands[i].resize(band_shapes[i], 0.0f); if (active0[i]) dequantize_band(quant0[i], steps0[i], ch0_bands[i]); - else { ch0_bands[i].clear(); ch0_bands[i].assign(band_shapes[i], 0.0f); } + if (stereo) { + ch1_bands[i].clear(); + ch1_bands[i].resize(band_shapes[i], 0.0f); if (active1[i]) dequantize_band(quant1[i], steps1[i], ch1_bands[i]); - else { ch1_bands[i].clear(); ch1_bands[i].assign(band_shapes[i], 0.0f); } } } - sbr_synthesize(ch0_bands, active0, band_shapes, sbr, sbr_end, bi); + sbr_synthesize(ch0_bands, active0, band_shapes, sbr, sbr_end, bi, SBR_W0, SBR_W1); if (stereo) { int is_start = get_is_start_band(target_kbps, total_bands); @@ -381,17 +623,20 @@ void decompress_wha_to_wav(const std::string& in_wha, if (block_format_version >= 18 && is_used && i >= is_start && is_start < expected_band_count) { std::vector left, right; apply_is(ch0_bands[i], is_r[i], is_inv[i], is_use_segmented[i], left, right); - ch0_bands[i] = std::move(left); ch1_bands[i] = std::move(right); + ch0_bands[i] = std::move(left); + ch1_bands[i] = std::move(right); } else { std::vector left_tmp, right_tmp; ms_to_lr(ch0_bands[i], ch1_bands[i], left_tmp, right_tmp); - ch0_bands[i] = std::move(left_tmp); ch1_bands[i] = std::move(right_tmp); + ch0_bands[i] = std::move(left_tmp); + ch1_bands[i] = std::move(right_tmp); } } else { if (block_format_version >= 18 && is_used && i >= is_start && is_start < expected_band_count) { std::vector left_tmp, right_tmp; ms_to_lr(ch0_bands[i], ch1_bands[i], left_tmp, right_tmp); - ch0_bands[i] = std::move(left_tmp); ch1_bands[i] = std::move(right_tmp); + ch0_bands[i] = std::move(left_tmp); + ch1_bands[i] = std::move(right_tmp); } } } @@ -413,16 +658,17 @@ void decompress_wha_to_wav(const std::string& in_wha, extract_ready_samples(ola, write_buf); compact_ola_buffer(ola); - if (write_buf.size() >= write_batch_frames * num_channels) { + if (write_buf.size() >= write_batch_frames * num_channels) flush_write_buffer(); - } if (pred_enabled) { pred_state.prev_active0.assign(expected_band_count, 0); - for (int i = 0; i < expected_band_count; ++i) if (active0[i]) pred_state.prev_active0[i] = 1; + for (int i = 0; i < expected_band_count; ++i) + if (active0[i]) pred_state.prev_active0[i] = 1; if (stereo) { pred_state.prev_active1.assign(expected_band_count, 0); - for (int i = 0; i < expected_band_count; ++i) if (active1[i]) pred_state.prev_active1[i] = 1; + for (int i = 0; i < expected_band_count; ++i) + if (active1[i]) pred_state.prev_active1[i] = 1; } else { pred_state.prev_active1.assign(expected_band_count, 0); } @@ -433,10 +679,12 @@ void decompress_wha_to_wav(const std::string& in_wha, if (idx >= 0 && idx < expected_band_count) { pred_state.prev_sbr_valid[idx] = 1; pred_state.prev_sbr_noise[idx] = sbr.noise_flag[idx] ? 1 : 0; - pred_state.prev_sbr_rms[idx] = sbr.rms_idx[idx]; + pred_state.prev_sbr_rms[idx] = sbr.rms_idx[idx]; } } - pred_state.ready = true; pred_state.last_bi = bi; + + pred_state.ready = true; + pred_state.last_bi = bi; } prev_block_size = block_size; @@ -445,17 +693,18 @@ void decompress_wha_to_wav(const std::string& in_wha, extract_remaining(ola, write_buf); flush_write_buffer(); - std::vector pow_tail; - pow_dec.flush(pow_tail); - write_samples(pow_tail); + if (!speed_mode && pow_dec) { + std::vector pow_tail; + pow_dec->flush(pow_tail); + write_samples(pow_tail); + } if (postproc) { int n_out = postproc->flush(); if (n_out > 0) { size_t frames_to_write = n_out; - if (total_frames > 0 && frames_written + frames_to_write > total_frames) { + if (total_frames > 0 && frames_written + frames_to_write > total_frames) frames_to_write = total_frames - frames_written; - } if (frames_to_write > 0) { drwav_write_pcm_frames(&wav, (uint64_t)frames_to_write, postproc->out_data()); frames_written += frames_to_write; @@ -464,17 +713,18 @@ void decompress_wha_to_wav(const std::string& in_wha, } drwav_uninit(&wav); + std::cout << "Decompressed: " << out_wav << " (frames=" << ola.output_pos << ", sr=" << sr << ", ch=" << (int)num_channels << ")" << std::endl; } -#include - -int main(int argc, char** argv) { +int main(int argc, char** argv) +{ if (argc < 3) { - std::cerr << "Usage: decoder [buffer_size_in_samples]" << std::endl; + std::cerr << "Usage: decoder [buffer_size_in_samples] [sbr_w0] [sbr_w1]\n"; + std::cerr << "Note: --speed flag is NO LONGER NEEDED. The decoder reads the mode from the container header automatically.\n"; return 1; } @@ -483,15 +733,22 @@ int main(int argc, char** argv) { try { buffer_size = std::stoul(argv[3]); if (buffer_size == 0) { - std::cerr << "Buffer size must be > 0, using default 8192" << std::endl; + std::cerr << "Buffer size must be > 0\n"; buffer_size = 8192; } } catch (...) { - std::cerr << "Invalid buffer size, using default 8192" << std::endl; + std::cerr << "Invalid buffer size\n"; buffer_size = 8192; } } + if (argc >= 5) { try { SBR_W0 = std::stof(argv[4]); } catch (...) {} } + if (argc >= 6) { try { SBR_W1 = std::stof(argv[5]); } catch (...) {} } + + std::cout << "SBR Shape Weights: W0=" << SBR_W0 + << ", W1=" << SBR_W1 + << " (Contrast: " << (20.0f * std::log10(SBR_W1 / SBR_W0)) << " dB)\n"; + try { struct rusage usage_before, usage_after; getrusage(RUSAGE_SELF, &usage_before); @@ -500,15 +757,15 @@ int main(int argc, char** argv) { getrusage(RUSAGE_SELF, &usage_after); double user_time_sec = - (usage_after.ru_utime.tv_sec - usage_before.ru_utime.tv_sec) + + (usage_after.ru_utime.tv_sec - usage_before.ru_utime.tv_sec) + (usage_after.ru_utime.tv_usec - usage_before.ru_utime.tv_usec) / 1000000.0; drwav wav_info; if (drwav_init_file(&wav_info, argv[2], nullptr)) { double duration = (double)wav_info.totalPCMFrameCount / (double)wav_info.sampleRate; double speed = (user_time_sec > 0.0) ? (duration / user_time_sec) : 0.0; - std::cout << "Decoding speed: " << speed << "x realtime (" - << user_time_sec << "s user time)" << std::endl; + std::cout << "Decoding speed: " << speed + << "x realtime (" << user_time_sec << "s user time)\n"; drwav_uninit(&wav_info); } } catch (const std::exception& ex) { diff --git a/decoder/joint_stereo.h b/decoder/joint_stereo.h index b7c15f4..4a41ff3 100644 --- a/decoder/joint_stereo.h +++ b/decoder/joint_stereo.h @@ -1,6 +1,5 @@ #ifndef JOINT_STEREO_H #define JOINT_STEREO_H - #include #include #include @@ -54,30 +53,21 @@ inline float compute_r_segment(const std::vector& left, const std::vector int start, int end) { double eL = 0.0, eR = 0.0; - for (int i = start; i < end; ++i) { eL += left[i] * left[i]; eR += right[i] * right[i]; } - const double eps = 1e-12; double total = eL + eR; - - if (total < eps) - return 0.5f; - + if (total < eps) return 0.5f; return static_cast(eR / total); } inline bool should_use_adaptive_panorama(const std::array& r_segments) { int changes = 0; - for (int i = 1; i < 4; ++i) { - if (std::abs(r_segments[i] - r_segments[i - 1]) > 0.3f) { - changes++; - } + if (std::abs(r_segments[i] - r_segments[i - 1]) > 0.3f) changes++; } - return (changes >= 2); } @@ -85,91 +75,78 @@ inline void compute_is_parameters_ex(const std::vector& left, const std::vector& right, std::vector& Y, std::array& r_out, - bool& inv_flag, + std::array& inv_flags, bool& use_segmented, int band_idx, int total_bands) { const size_t n = left.size(); - double eL_total = 0.0, eR_total = 0.0, dot_total = 0.0; - for (size_t i = 0; i < n; ++i) { eL_total += left[i] * left[i]; eR_total += right[i] * right[i]; dot_total += left[i] * right[i]; } + + for (int s = 0; s < 4; ++s) inv_flags[s] = false; const double eps = 1e-12; double total = eL_total + eR_total; - if (total < eps || eL_total < eps || eR_total < eps) { - for (int s = 0; s < 4; ++s) - r_out[s] = 0.5f; - + for (int s = 0; s < 4; ++s) r_out[s] = 0.5f; Y.assign(n, 0.0f); - inv_flag = false; use_segmented = false; return; } - double corr = dot_total / (std::sqrt(eL_total * eR_total) + eps); - inv_flag = (corr < -0.2); - - int segmented_threshold = 6 * total_bands / 16; + int segmented_threshold = 8 * total_bands / 16; use_segmented = (band_idx < segmented_threshold); if (use_segmented) { const int seg_size = (int)(n / 4); Y.resize(n); - for (int s = 0; s < 4; ++s) { int start = s * seg_size; int end = (s == 3) ? (int)n : start + seg_size; int len = end - start; - - double seg_eL = 0.0, seg_eR = 0.0; - + + double seg_eL = 0.0, seg_eR = 0.0, seg_dot = 0.0; for (int i = start; i < end; ++i) { seg_eL += left[i] * left[i]; seg_eR += right[i] * right[i]; + seg_dot += left[i] * right[i]; } - + double seg_total = seg_eL + seg_eR; - if (seg_total < eps) { r_out[s] = 0.5f; - for (int i = start; i < end; ++i) - Y[i] = 0.0f; + for (int i = start; i < end; ++i) Y[i] = 0.0f; continue; } - + + double seg_corr = seg_dot / (std::sqrt(seg_eL * seg_eR) + eps); + inv_flags[s] = (seg_corr < -0.2); + r_out[s] = static_cast(seg_eR / seg_total); - double sqrt_eL = std::sqrt(seg_eL); double sqrt_eR = std::sqrt(seg_eR); double inv_sqrt_total = 1.0 / std::sqrt(seg_total); - + double eX = 0.0; std::vector X_seg(len); - for (int i = 0; i < len; ++i) { int idx = start + i; - - double Rterm = inv_flag ? -sqrt_eR * right[idx] : sqrt_eR * right[idx]; + double Rterm = inv_flags[s] ? -sqrt_eR * right[idx] : sqrt_eR * right[idx]; double x = (sqrt_eL * left[idx] + Rterm) * inv_sqrt_total; - X_seg[i] = static_cast(x); eX += x * x; } - + if (eX < eps) { - for (int i = start; i < end; ++i) - Y[i] = 0.0f; + for (int i = start; i < end; ++i) Y[i] = 0.0f; } else { double inv_sqrt_eX = 1.0 / std::sqrt(eX); double target_energy = std::sqrt(seg_total); - for (int i = 0; i < len; ++i) { Y[start + i] = static_cast(X_seg[i] * inv_sqrt_eX * target_energy); } @@ -177,33 +154,31 @@ inline void compute_is_parameters_ex(const std::vector& left, } } else { use_segmented = false; - + double corr = dot_total / (std::sqrt(eL_total * eR_total) + eps); + inv_flags[0] = (corr < -0.2); + float r = static_cast(eR_total / total); - for (int s = 0; s < 4; ++s) - r_out[s] = r; - + for (int s = 0; s < 4; ++s) r_out[s] = r; + double sqrtE = std::sqrt(total); double sqrt_eL = std::sqrt(eL_total); double sqrt_eR = std::sqrt(eR_total); double inv_sqrt_total = 1.0 / std::sqrt(total); - + std::vector X(n); double eX = 0.0; - for (size_t i = 0; i < n; ++i) { - double Rterm = inv_flag ? -sqrt_eR * right[i] : sqrt_eR * right[i]; + double Rterm = inv_flags[0] ? -sqrt_eR * right[i] : sqrt_eR * right[i]; double x = (sqrt_eL * left[i] + Rterm) * inv_sqrt_total; - X[i] = static_cast(x); eX += x * x; } - + if (eX < eps) { Y.assign(n, 0.0f); } else { double inv_sqrt_eX = 1.0 / std::sqrt(eX); Y.resize(n); - for (size_t i = 0; i < n; ++i) { Y[i] = static_cast(X[i] * inv_sqrt_eX * sqrtE); } @@ -213,7 +188,7 @@ inline void compute_is_parameters_ex(const std::vector& left, inline void apply_is(const std::vector& Y, const std::array& r_vals, - bool inv_flag, + const std::array& inv_flags, bool use_segmented, std::vector& left, std::vector& right) @@ -221,31 +196,30 @@ inline void apply_is(const std::vector& Y, size_t n = Y.size(); left.resize(n); right.resize(n); - + if (use_segmented && n >= 16) { int seg_size = (int)(n / 4); - for (int s = 0; s < 4; ++s) { float r = r_vals[s]; float gl = std::sqrt(1.0f - r); float gr = std::sqrt(r); - int start = s * seg_size; int end = (s == 3) ? (int)n : start + seg_size; - + + bool local_inv = inv_flags[s]; for (int i = start; i < end; ++i) { left[i] = Y[i] * gl; - right[i] = Y[i] * (inv_flag ? -gr : gr); + right[i] = Y[i] * (local_inv ? -gr : gr); } } } else { float r = r_vals[0]; float gl = std::sqrt(1.0f - r); float gr = std::sqrt(r); - + bool local_inv = inv_flags[0]; for (size_t i = 0; i < n; ++i) { left[i] = Y[i] * gl; - right[i] = Y[i] * (inv_flag ? -gr : gr); + right[i] = Y[i] * (local_inv ? -gr : gr); } } } @@ -253,20 +227,16 @@ inline void apply_is(const std::vector& Y, inline uint32_t quantize_r(float r, int bits) { float r_clamped = std::clamp(r, 0.0f, 1.0f); float theta = std::asin(std::sqrt(r_clamped)); - int max_val = (1 << bits) - 2; float theta_norm = theta / (float)(M_PI * 0.5f); - return static_cast(std::clamp(theta_norm, 0.0f, 1.0f) * max_val + 0.5f); } inline float dequantize_r(uint32_t q, int bits) { int max_val = (1 << bits) - 2; - float theta_norm = (float)q / max_val; float theta = theta_norm * (float)(M_PI * 0.5f); float s = std::sin(theta); - return s * s; } @@ -288,37 +258,22 @@ inline std::pair, std::vector> mid_side(const std::vec inline bool use_mid_side(float El, float Er, float Em, float Es, bool enable_ms, float target_kbps) { - if (!enable_ms) - return false; - + if (!enable_ms) return false; const float eps = 1e-12f; - - if (target_kbps < 128.0f) { - if (Es > Em * 0.9f) - return false; - + if (target_kbps < 64.0f) { + if (Es > Em * 0.9f) return false; float prod_lr = El * Er; float prod_ms = Em * Es; - - if (prod_ms > prod_lr * 1.05f) - return false; - + if (prod_ms > prod_lr * 1.05f) return false; return true; } - float min_e = std::min(El, Er); float max_e = std::max(El, Er); float energy_ratio = max_e / (min_e + eps); - - if (energy_ratio > 4.0f) - return false; - - if (Es > Em * 0.5f) - return false; - + if (energy_ratio > 4.0f) return false; + if (Es > Em * 0.5f) return false; float prod_lr = El * Er; float prod_ms = Em * Es; - return prod_ms < prod_lr * 0.95f; } @@ -330,49 +285,57 @@ inline bool should_use_is_band(const std::vector& left, float equal_ratio_threshold, float equal_ratio_threshold_low_bitrate) { - size_t n = std::min(left.size(), right.size()); - if (n == 0) - return false; - - double eL = 0.0, eR = 0.0; - + const size_t n = std::min(left.size(), right.size()); + if (n == 0) return false; + const double eps = 1e-12; + double eL = 0.0, eR = 0.0, cross = 0.0; for (size_t i = 0; i < n; ++i) { - eL += (double)left[i] * left[i]; - eR += (double)right[i] * right[i]; + const double L = (double)left[i]; + const double R = (double)right[i]; + eL += L * L; eR += R * R; cross += L * R; } - - const double eps = 1e-12; - if (eL + eR < eps) - return false; - - double meanL = eL / (double)n; - double meanR = eR / (double)n; - - double minE = std::min(meanL, meanR); - double maxE = std::max(meanL, meanR); - double ratio = maxE / (minE + eps); - - if (ratio > (double)strong_ratio_threshold) - return true; - + if (eL + eR < eps) return false; + const double correlation = cross / (std::sqrt(eL * eR) + eps); + const double abs_correlation = std::abs(correlation); + constexpr double IS_CORRELATION_THRESHOLD = 0.40; + if (abs_correlation >= IS_CORRELATION_THRESHOLD) return true; + + const double meanL = eL / (double)n; + const double meanR = eR / (double)n; + const double minE = std::min(meanL, meanR); + const double maxE = std::max(meanL, meanR); + const double ratio = maxE / (minE + eps); + if (ratio > (double)strong_ratio_threshold) return true; + double eM = 0.0, eS = 0.0; - for (size_t i = 0; i < n; ++i) { - double m = ((double)left[i] + (double)right[i]) * 0.5; - double s = ((double)left[i] - (double)right[i]) * 0.5; - - eM += m * m; - eS += s * s; + const double L = (double)left[i]; + const double R = (double)right[i]; + const double M = (L + R) * 0.5; + const double S = (L - R) * 0.5; + eM += M * M; eS += S * S; } + const bool ms_benefit = use_mid_side((float)eL, (float)eR, (float)eM, (float)eS, enable_ms, target_kbps); + return ms_benefit || maxE >= minE * 4.0; +} - bool ms_benefit = use_mid_side((float)eL, (float)eR, (float)eM, (float)eS, - enable_ms, target_kbps); - - float equal_thr = (target_kbps < 128.0f) - ? equal_ratio_threshold_low_bitrate - : equal_ratio_threshold; +inline int32_t quantize_r_centered(float r, int bits) { + float r_clamped = std::clamp(r, 0.0f, 1.0f); + float theta = std::asin(std::sqrt(r_clamped)); + float theta_center = (float)(M_PI * 0.25); + int steps = (1 << (bits - 1)) - 1; + float step_size = theta_center / steps; + int32_t q = (int32_t)std::round((theta - theta_center) / step_size); + return std::clamp(q, -steps, steps); +} - return (ratio <= (double)equal_thr && ms_benefit); +inline float dequantize_r_centered(int32_t q, int bits) { + int steps = (1 << (bits - 1)) - 1; + float theta_center = (float)(M_PI * 0.25); + float step_size = theta_center / steps; + float theta = theta_center + q * step_size; + float s = std::sin(theta); + return s * s; } #endif // JOINT_STEREO_H diff --git a/decoder/pred_state.h b/decoder/pred_state.h index ccd6108..029f0dd 100644 --- a/decoder/pred_state.h +++ b/decoder/pred_state.h @@ -8,6 +8,9 @@ struct ModeState { bool ready = false; uint32_t last_bi = 0; + std::vector prev_k_sub0; + std::vector prev_k_sub1; + std::vector prev_scale0; std::vector prev_scale1; std::vector prev_active0; diff --git a/decoder/sbr_decode.h b/decoder/sbr_decode.h index d012be6..52c8c12 100644 --- a/decoder/sbr_decode.h +++ b/decoder/sbr_decode.h @@ -8,14 +8,98 @@ #include "pred_state.h" #include "entropy_decoder.h" +extern float SBR_W0; +extern float SBR_W1; + +inline void get_sbr_shape_weights(uint8_t mask, float w0, float w1, float out[4]) { + float sum = 0.0f; + float v[4]; + for (int i = 0; i < 4; ++i) { + v[i] = (mask & (1 << i)) ? w1 : w0; + sum += v[i]; + } + for (int i = 0; i < 4; ++i) { + out[i] = v[i] / sum; + } +} + struct SBRDecodeResult { std::vector noise_flag; std::vector rms; std::vector rms_idx; + std::vector shape_mask; std::vector indices; bool used = false; }; +inline SBRDecodeResult sbr_decode_v23(BitReaderMSB& reader, + const std::vector& active0, + int sbr_end, int expected_bands, + bool can_state_pred, + const ModeState& pred_state) +{ + SBRDecodeResult result; + result.noise_flag.assign(expected_bands, false); + result.rms.assign(expected_bands, 0.0f); + result.rms_idx.assign(expected_bands, 0); + result.shape_mask.assign(expected_bands, 0); + + if (sbr_end <= 0) return result; + + std::vector indices; + for (int i = 0; i < sbr_end; ++i) { + if (!active0[i]) indices.push_back(i); + } + if (indices.empty()) return result; + + result.indices = indices; + result.used = true; + + for (int idx : indices) { + bool bit = (reader.read_bits(1) != 0); + if (can_state_pred && idx < (int)pred_state.prev_sbr_valid.size() && pred_state.prev_sbr_valid[idx]) { + bit = bit ^ (pred_state.prev_sbr_noise[idx] != 0); + } + result.noise_flag[idx] = bit; + } + + for (int idx : indices) { + result.shape_mask[idx] = (uint8_t)reader.read_bits(4); + } + + int k_sbr = (int)reader.read_bits(3); + auto sbr_vals = rice_decode(reader, (int)indices.size(), k_sbr); + + if ((int)sbr_vals.size() != (int)indices.size()) { + throw std::runtime_error("SBR RMS decode size mismatch"); + } + + for (size_t j = 0; j < indices.size(); ++j) { + int idx = indices[j]; + uint32_t val = sbr_vals[j]; + int sb = get_scale_bits(idx); + int max_idx = (1 << sb) - 1; + uint32_t rms_idx; + + if (can_state_pred && idx < (int)pred_state.prev_sbr_valid.size() && pred_state.prev_sbr_valid[idx]) { + int32_t d = zigzag_decode(val); + int64_t v = (int64_t)pred_state.prev_sbr_rms[idx] + d; + if (v < 0) v = 0; + if (v > max_idx) v = max_idx; + rms_idx = (uint32_t)v; + } else { + rms_idx = val; + if (rms_idx > (uint32_t)max_idx) rms_idx = max_idx; + } + + result.rms_idx[idx] = rms_idx; + float log_rms = SCALE_LOG_MIN + rms_idx * (SCALE_LOG_MAX - SCALE_LOG_MIN) / max_idx; + result.rms[idx] = exp10f(log_rms); + } + + return result; +} + inline SBRDecodeResult sbr_decode(const uint8_t* blk, size_t& ptr, size_t blk_size, const std::vector& active0, int sbr_end, int expected_bands, @@ -27,35 +111,38 @@ inline SBRDecodeResult sbr_decode(const uint8_t* blk, size_t& ptr, size_t blk_si result.noise_flag.assign(expected_bands, false); result.rms.assign(expected_bands, 0.0f); result.rms_idx.assign(expected_bands, 0); - + result.shape_mask.assign(expected_bands, 0); if (sbr_end <= 0) return result; - for (int i = 0; i < sbr_end; ++i) { - if (!active0[i]) - result.indices.push_back(i); + if (!active0[i]) result.indices.push_back(i); } if (result.indices.empty()) return result; int noise_bytes = ((int)result.indices.size() + 7) / 8; if (ptr + noise_bytes > blk_size) return result; - for (size_t j = 0; j < result.indices.size(); ++j) { int idx = result.indices[j]; bool bit = (blk[ptr + j / 8] >> (j % 8)) & 1; - if (can_state_pred && - idx < (int)pred_state.prev_sbr_valid.size() && - pred_state.prev_sbr_valid[idx]) - { + if (can_state_pred && idx < (int)pred_state.prev_sbr_valid.size() && pred_state.prev_sbr_valid[idx]) { bit = bit ^ (pred_state.prev_sbr_noise[idx] != 0); } result.noise_flag[idx] = bit; } ptr += noise_bytes; + int mask_count = result.indices.size(); + int mask_bytes = (mask_count + 1) / 2; + if (ptr + mask_bytes > blk_size) return result; + for (int j = 0; j < mask_count; ++j) { + uint8_t byte = blk[ptr + j / 2]; + uint8_t m = (j % 2 == 0) ? (byte & 0x0F) : ((byte >> 4) & 0x0F); + result.shape_mask[result.indices[j]] = m; + } + ptr += mask_bytes; + if (block_format_version >= 19) { if (ptr >= blk_size) return result; uint8_t sbr_rms_mode = blk[ptr++]; - if (sbr_rms_mode == 1) { if (ptr >= blk_size) return result; int k_sbr = blk[ptr++]; @@ -63,27 +150,18 @@ inline SBRDecodeResult sbr_decode(const uint8_t* blk, size_t& ptr, size_t blk_si uint16_t rms_len; memcpy(&rms_len, &blk[ptr], 2); ptr += 2; - if (ptr + rms_len > blk_size) - throw std::runtime_error("SBR RMS truncated"); - + if (ptr + rms_len > blk_size) throw std::runtime_error("SBR RMS truncated"); BitReaderMSB rms_reader(blk + ptr, rms_len); ptr += rms_len; result.used = true; - for (int idx : result.indices) { auto val_vec = rice_decode(rms_reader, 1, k_sbr); - if (val_vec.empty()) - throw std::runtime_error("SBR RMS index missing"); + if (val_vec.empty()) throw std::runtime_error("SBR RMS index missing"); uint32_t val = val_vec[0]; - - int sb = get_scale_bits(idx); + int sb = get_scale_bits(idx); int max_idx = (1 << sb) - 1; uint32_t rms_idx; - - if (can_state_pred && - idx < (int)pred_state.prev_sbr_valid.size() && - pred_state.prev_sbr_valid[idx]) - { + if (can_state_pred && idx < (int)pred_state.prev_sbr_valid.size() && pred_state.prev_sbr_valid[idx]) { int32_t d = zigzag_decode(val); int64_t v = (int64_t)pred_state.prev_sbr_rms[idx] + d; if (v < 0) v = 0; @@ -91,62 +169,48 @@ inline SBRDecodeResult sbr_decode(const uint8_t* blk, size_t& ptr, size_t blk_si rms_idx = (uint32_t)v; } else { rms_idx = val; - if (rms_idx > (uint32_t)max_idx) - rms_idx = max_idx; + if (rms_idx > (uint32_t)max_idx) rms_idx = max_idx; } - result.rms_idx[idx] = rms_idx; - float log_rms = SCALE_LOG_MIN + - rms_idx * (SCALE_LOG_MAX - SCALE_LOG_MIN) / max_idx; + float log_rms = SCALE_LOG_MIN + rms_idx * (SCALE_LOG_MAX - SCALE_LOG_MIN) / max_idx; result.rms[idx] = exp10f(log_rms); } } else { int total_rms_bits = 0; - for (int idx : result.indices) - total_rms_bits += get_scale_bits(idx); + for (int idx : result.indices) total_rms_bits += get_scale_bits(idx); int rms_bytes = (total_rms_bits + 7) / 8; if (ptr + rms_bytes > blk_size) return result; - BitReaderMSB rms_reader(blk + ptr, rms_bytes); result.used = true; - for (int idx : result.indices) { - int sb = get_scale_bits(idx); + int sb = get_scale_bits(idx); uint32_t rms_idx = rms_reader.read_bits(sb); - int max_idx = (1 << sb) - 1; - if (rms_idx > (uint32_t)max_idx) - rms_idx = max_idx; + int max_idx = (1 << sb) - 1; + if (rms_idx > (uint32_t)max_idx) rms_idx = max_idx; result.rms_idx[idx] = rms_idx; - float log_rms = SCALE_LOG_MIN + - rms_idx * (SCALE_LOG_MAX - SCALE_LOG_MIN) / max_idx; + float log_rms = SCALE_LOG_MIN + rms_idx * (SCALE_LOG_MAX - SCALE_LOG_MIN) / max_idx; result.rms[idx] = exp10f(log_rms); } ptr += rms_bytes; } } else { int total_rms_bits = 0; - for (int idx : result.indices) - total_rms_bits += get_scale_bits(idx); + for (int idx : result.indices) total_rms_bits += get_scale_bits(idx); int rms_bytes = (total_rms_bits + 7) / 8; if (ptr + rms_bytes > blk_size) return result; - BitReaderMSB rms_reader(blk + ptr, rms_bytes); result.used = true; - for (int idx : result.indices) { - int sb = get_scale_bits(idx); + int sb = get_scale_bits(idx); uint32_t rms_idx = rms_reader.read_bits(sb); - int max_idx = (1 << sb) - 1; - if (rms_idx > (uint32_t)max_idx) - rms_idx = max_idx; + int max_idx = (1 << sb) - 1; + if (rms_idx > (uint32_t)max_idx) rms_idx = max_idx; result.rms_idx[idx] = rms_idx; - float log_rms = SCALE_LOG_MIN + - rms_idx * (SCALE_LOG_MAX - SCALE_LOG_MIN) / max_idx; + float log_rms = SCALE_LOG_MIN + rms_idx * (SCALE_LOG_MAX - SCALE_LOG_MIN) / max_idx; result.rms[idx] = exp10f(log_rms); } ptr += rms_bytes; } - return result; } @@ -156,156 +220,83 @@ inline void sbr_synthesize( const std::vector& band_shapes, const SBRDecodeResult& sbr, int sbr_end, - uint32_t block_idx) + uint32_t block_idx, + float w0, float w1) { - if (!sbr.used) - return; - + if (!sbr.used) return; const int band_count = (int)ch0_bands.size(); - std::vector old_index(band_count); - - for (int i = 0; i < band_count; ++i) - old_index[i] = i ^ (i >> 1); - + for (int i = 0; i < band_count; ++i) old_index[i] = i ^ (i >> 1); + std::vector available_sources; - for (int j = 0; j < band_count; ++j) { - if (active0[j]) - available_sources.push_back(j); + if (active0[j]) available_sources.push_back(j); } - - std::vector source_used( - band_count, - false - ); - + + std::vector source_used(band_count, false); for (int i = 0; i < sbr_end; ++i) { - - if (active0[i]) - continue; - + if (active0[i]) continue; const int N = band_shapes[i]; - - const float target_rms = - sbr.rms[i]; - - const float target_energy = - target_rms * target_rms * N; - + const float target_rms = sbr.rms[i]; + const float target_energy = target_rms * target_rms * N; ch0_bands[i].resize(N); - + if (sbr.noise_flag[i]) { - - uint32_t seed = - block_idx * 1000000u + - (uint32_t)i * 137u + - 1u; - - float noise_energy = 0.0f; - + uint32_t seed = block_idx * 1000000u + (uint32_t)i * 137u + 1u; for (int j = 0; j < N; ++j) { - - seed = - seed * 1664525u + - 1013904223u; - - float noise = - (static_cast(seed) / - 2147483648.0f) - 1.0f; - - ch0_bands[i][j] = noise; - - noise_energy += - noise * noise; + seed = seed * 1664525u + 1013904223u; + ch0_bands[i][j] = (static_cast(seed) / 2147483648.0f) - 1.0f; } - - const float gain = - std::sqrt( - target_energy / - (noise_energy + 1e-12f) - ); - - for (int j = 0; j < N; ++j) - ch0_bands[i][j] *= gain; - } - - else { - + } else { int source_band = -1; float best_weight = -1.0f; - const int IDEAL_DISTANCE = 2; const float SIGMA = 3.0f; - - const int target_old = - old_index[i]; - + const int target_old = old_index[i]; for (int src : available_sources) { - - if (source_used[src]) - continue; - - const int source_old = - old_index[src]; - - const int distance = - std::abs( - source_old - target_old - ); - - const float d = - (float)(distance - - IDEAL_DISTANCE); - - const float weight = - std::exp( - -(d * d) / - (2.0f * - SIGMA * - SIGMA) - ); - + if (source_used[src]) continue; + const int distance = std::abs((old_index[src]) - target_old); + const float d = (float)(distance - IDEAL_DISTANCE); + const float weight = std::exp(-(d * d) / (2.0f * SIGMA * SIGMA)); if (weight > best_weight) { best_weight = weight; source_band = src; } } - if (source_band >= 0) { - source_used[source_band] = true; - - const auto& src = - ch0_bands[source_band]; - - const int src_size = - (int)src.size(); - - float src_energy = 0.0f; - - for (float v : src) - src_energy += v * v; - - const float gain = - std::sqrt( - target_energy / - (src_energy + 1e-12f) - ); - - for (int j = 0; j < N; ++j) - ch0_bands[i][j] = - src[j % src_size] * gain; + const auto& src = ch0_bands[source_band]; + const int src_size = (int)src.size(); + for (int j = 0; j < N; ++j) { + ch0_bands[i][j] = src[j % src_size]; + } + } else { + std::fill(ch0_bands[i].begin(), ch0_bands[i].end(), 0.0f); } - - else { - std::fill( - ch0_bands[i].begin(), - ch0_bands[i].end(), - 0.0f - ); + } + + int seg_size = N / 4; + float seg_src_energy[4] = {0, 0, 0, 0}; + for (int s = 0; s < 4; ++s) { + int start = s * seg_size; + int end = (s == 3) ? N : (s + 1) * seg_size; + for (int j = start; j < end; ++j) { + seg_src_energy[s] += ch0_bands[i][j] * ch0_bands[i][j]; + } + } + float weights[4]; + get_sbr_shape_weights(sbr.shape_mask[i], w0, w1, weights); + for (int s = 0; s < 4; ++s) { + int start = s * seg_size; + int end = (s == 3) ? N : (s + 1) * seg_size; + float seg_target_e = target_energy * weights[s]; + float gain = std::sqrt(seg_target_e / (seg_src_energy[s] + 1e-12f)); + if (gain > 2) { + gain = 0; + } + for (int j = start; j < end; ++j) { + ch0_bands[i][j] *= gain; } } } } -