Spaces:
Sleeping
Sleeping
| import torch | |
| def compute_rms(x: torch.Tensor, **kwargs): | |
| """Compute root mean square energy. | |
| Args: | |
| x: (bs, 1, seq_len) | |
| Returns: | |
| rms: (bs, ) | |
| """ | |
| rms = torch.sqrt(torch.mean(x**2, dim=-1).clamp(min=1e-8)) | |
| return rms | |
| def compute_log_rms_gated_max(x: torch.Tensor, sample_rate=44100, **kwargs): | |
| """Compute gated log RMS energy. | |
| Frames the signal in 400 ms windows with 75% overlap, computes RMS, | |
| discards frames with RMS < -60 dBFS, and averages the log-RMS. | |
| If all frames in a given (batch, channel) are below -60 dBFS, | |
| returns -60 for that entry. | |
| Args: | |
| x: Tensor of shape (bs, c, seq_len) | |
| Returns: | |
| log_rms: Tensor of shape (bs, c) | |
| """ | |
| seg_size = int(sample_rate * 0.4) | |
| hop_size = int(sample_rate * 0.1) | |
| # (bs, c, num_frames, seg_size) | |
| B, C, L = x.size() | |
| assert C==1 or C==2 | |
| x_frames = x.unfold(2, seg_size, hop_size) # (bs, c, num_frames, seg_size) | |
| # RMS over last dimension (seg_size) | |
| rms = torch.sqrt((x_frames ** 2).mean(dim=-1)) # (bs, c, num_frames) | |
| # dB conversion | |
| rms_db = 20 * torch.log10(rms.clamp(min=1e-8)) # (bs, c, num_frames) | |
| #print("rms db shape", rms_db.shape) | |
| #take the maximum RMS across all frames | |
| rms_max = rms_db.max(dim=2)[0] # (bs, c) | |
| #print(f"RMS max shape: {rms_max.shape}") | |
| return rms_max | |
| def compute_log_rms(x: torch.Tensor, **kwargs): | |
| """Compute root mean square energy. | |
| Args: | |
| x: (bs, 1, seq_len) | |
| Returns: | |
| rms: (bs, ) | |
| """ | |
| rms=compute_rms(x) | |
| return 20 * torch.log10(rms.clamp(min=1e-8)) | |
| def compute_crest_factor(x: torch.Tensor, **kwargs): | |
| """Compute crest factor as ratio of peak to rms energy in dB. | |
| Args: | |
| x: (bs, 2, seq_len) | |
| """ | |
| num = torch.max(torch.abs(x), dim=-1)[0] | |
| den = compute_rms(x).clamp(min=1e-8) | |
| cf = 20 * torch.log10((num / den).clamp(min=1e-8)) | |
| return cf | |
| def compute_log_spread(x: torch.Tensor, **kwargs): | |
| """Compute log spread as the mean difference between log magnitude of samples and log RMS. | |
| Args: | |
| x: (bs, 1, seq_len) | |
| Returns: | |
| log_spread: (bs, ) | |
| """ | |
| # Compute log RMS | |
| log_rms = compute_log_rms(x).unsqueeze(-1) # (bs, 1, 1) | |
| # Compute log magnitude of each sample | |
| log_magnitude = 20 * torch.log10(torch.abs(x).clamp(min=1e-8)) # (bs, 1, seq_len) | |
| # Compute the difference and take the mean | |
| log_spread = torch.mean(log_magnitude - log_rms, dim=-1).squeeze(1) # (bs, ) | |
| return log_spread | |
| def compute_stereo_width(x: torch.Tensor, **kwargs): | |
| """Compute stereo width as ratio of energy in sum and difference signals. | |
| Args: | |
| x: (bs, 2, seq_len) | |
| """ | |
| bs, chs, seq_len = x.size() | |
| assert chs == 2, "Input must be stereo" | |
| # compute sum and diff of stereo channels | |
| x_sum = x[:, 0, :] + x[:, 1, :] | |
| x_diff = x[:, 0, :] - x[:, 1, :] | |
| # compute power of sum and diff | |
| sum_energy = torch.mean(x_sum**2, dim=-1) | |
| diff_energy = torch.mean(x_diff**2, dim=-1) | |
| # compute stereo width as ratio | |
| stereo_width = diff_energy / sum_energy.clamp(min=1e-8) | |
| return stereo_width | |
| def compute_stereo_imbalance(x: torch.Tensor, **kwargs): | |
| """Compute stereo imbalance as ratio of energy in left and right channels. | |
| Args: | |
| x: (bs, 2, seq_len) | |
| Returns: | |
| stereo_imbalance: (bs, ) | |
| """ | |
| left_energy = torch.mean(x[:, 0, :] ** 2, dim=-1) | |
| right_energy = torch.mean(x[:, 1, :] ** 2, dim=-1) | |
| stereo_imbalance = (right_energy - left_energy) / ( | |
| right_energy + left_energy | |
| ).clamp(min=1e-8) | |
| return stereo_imbalance | |