144 lines
4.6 KiB
Python
144 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Author: C0ffymachyne
|
|
License: GPLv3
|
|
Version: 1.0.0
|
|
|
|
Description:
|
|
various normalization methods
|
|
"""
|
|
|
|
import torch
|
|
import pyloudnorm as pyln # pip install pyloudnorm
|
|
import numpy as np
|
|
|
|
|
|
def rms_normalization(audio: torch.Tensor, target_rms: float = 0.1) -> torch.Tensor:
|
|
rms = torch.sqrt(torch.mean(audio**2))
|
|
scaling_factor = target_rms / rms
|
|
normalized_audio = audio * scaling_factor
|
|
return normalized_audio
|
|
|
|
|
|
def lufs_normalization(
|
|
audio: torch.Tensor, sample_rate: int, target_lufs: float = -14.0
|
|
) -> torch.Tensor:
|
|
meter = pyln.Meter(sample_rate) # Create a loudness meter
|
|
|
|
__audio = audio.T.cpu().numpy()
|
|
loudness = meter.integrated_loudness(__audio) # Current LUFS
|
|
|
|
loudness_offset = target_lufs - loudness
|
|
normalized_audio = __audio * (10 ** (loudness_offset / 20.0))
|
|
|
|
result = torch.from_numpy(normalized_audio)
|
|
result = result.T
|
|
|
|
result = result.to(device=audio.device, dtype=result.dtype)
|
|
|
|
return result
|
|
|
|
|
|
def peak_normalization(audio: torch.Tensor, target_peak: float = 0.9) -> torch.Tensor:
|
|
peak = torch.max(torch.abs(audio))
|
|
scaling_factor = target_peak / peak
|
|
normalized_audio = audio * scaling_factor
|
|
return normalized_audio
|
|
|
|
|
|
def get_loudness(audio: torch.Tensor, sample_rate: int) -> float:
|
|
meter = pyln.Meter(sample_rate) # Create a loudness meter
|
|
audio = audio.T
|
|
audio = audio.cpu()
|
|
loudness: float = float(meter.integrated_loudness(audio.numpy())) # Current LUFS
|
|
return loudness
|
|
|
|
|
|
def set_loudness2(
|
|
audio_signal: torch.Tensor, sample_rate: int, target_loudness_db: float = -20.0
|
|
) -> torch.Tensor:
|
|
"""
|
|
Adjusts the loudness of the audio signal to a target level in decibels.
|
|
|
|
Args:
|
|
audio_signal (torch.Tensor): Input audio signal (channels, samples).
|
|
sample_rate (int): Sample rate of the audio signal.
|
|
target_loudness_db (float): Desired loudness in dB (e.g., -20.0 dB).
|
|
|
|
Returns:
|
|
torch.Tensor: Audio signal adjusted to the target loudness.
|
|
"""
|
|
# Convert PyTorch tensor to NumPy array for loudness calculation
|
|
audio_np = audio_signal.cpu().numpy().T # Convert to [samples, channels]
|
|
|
|
# Use pyloudnorm Meter to calculate and normalize loudness
|
|
meter = pyln.Meter(sample_rate) # Create loudness meter
|
|
current_loudness = meter.integrated_loudness(audio_np) # Measure LUFS
|
|
|
|
# Compute loudness adjustment gain
|
|
loudness_offset = target_loudness_db - current_loudness
|
|
gain_factor = 10 ** (loudness_offset / 20.0)
|
|
|
|
# Apply gain to adjust loudness
|
|
adjusted_audio_np = audio_np * gain_factor
|
|
|
|
# Convert back to PyTorch tensor
|
|
adjusted_audio = torch.from_numpy(adjusted_audio_np.T).to(
|
|
audio_signal.device, dtype=torch.float32
|
|
)
|
|
|
|
return adjusted_audio
|
|
|
|
|
|
def set_loudness(
|
|
audio_signal: torch.Tensor, sample_rate: int, target_loudness_db: float = -20.0
|
|
) -> torch.Tensor:
|
|
"""
|
|
Adjusts the loudness of the audio signal to a target level in decibels,
|
|
ensuring no clipping occurs.
|
|
|
|
Args:
|
|
audio_signal (torch.Tensor): Input audio signal (channels, samples).
|
|
sample_rate (int): Sample rate of the audio signal.
|
|
target_loudness_db (float): Desired loudness in dB (e.g., -20.0 dB).
|
|
|
|
Returns:
|
|
torch.Tensor: Audio signal adjusted to the target loudness.
|
|
"""
|
|
# Convert PyTorch tensor to NumPy array for loudness calculation
|
|
audio_np = audio_signal.cpu().numpy().T # Convert to [samples, channels]
|
|
|
|
# Use pyloudnorm Meter to calculate and normalize loudness
|
|
meter = pyln.Meter(sample_rate) # Create loudness meter
|
|
current_loudness = meter.integrated_loudness(audio_np) # Measure LUFS
|
|
|
|
# Compute loudness adjustment gain
|
|
loudness_offset = target_loudness_db - current_loudness
|
|
gain_factor = 10 ** (loudness_offset / 20.0)
|
|
|
|
# Apply gain to adjust loudness
|
|
adjusted_audio_np = audio_np * gain_factor
|
|
|
|
# Prevent clipping by normalizing the peak
|
|
peak_amplitude = np.max(np.abs(adjusted_audio_np))
|
|
if peak_amplitude > 1.0:
|
|
adjusted_audio_np = adjusted_audio_np / peak_amplitude
|
|
|
|
# Convert back to PyTorch tensor
|
|
adjusted_audio = torch.from_numpy(adjusted_audio_np.T).to(
|
|
audio_signal.device, dtype=torch.float32
|
|
)
|
|
|
|
return adjusted_audio
|
|
|
|
|
|
def automatic_gain_control(
|
|
audio: torch.Tensor, target_level: float = 0.7, alpha: float = 0.1
|
|
) -> torch.Tensor:
|
|
current_level = torch.mean(torch.abs(audio))
|
|
gain = target_level / (current_level + 1e-6)
|
|
smoothed_gain = alpha * gain + (1 - alpha) * 1.0
|
|
agc_audio = audio * smoothed_gain
|
|
return agc_audio
|