Upload 19 files

utils/bandwidth_sub.py

@@ -0,0 +1,123 @@
+import numpy as np
+import soundfile as sf
+import librosa
+import os
+from scipy.signal import butter, filtfilt, stft, istft
+
+# Step 1: Load audio files
+def load_audio(audio_path):
+    audio, sr = librosa.load(audio_path, sr=48000)
+    #audio, fs = sf.read(audio_path)
+    return audio, sr
+
+# Step 2: Detect effective signal bandwidth
+def detect_bandwidth_org(signal, fs, energy_threshold=0.95):
+    f, t, Zxx = stft(signal, fs=fs)
+    psd = np.abs(Zxx)**2
+    total_energy = np.sum(psd)
+    cumulative_energy = np.cumsum(np.sum(psd, axis=1)) / total_energy
+    f_low = f[np.argmax(cumulative_energy > (1 - energy_threshold))]
+    f_high = f[np.argmax(cumulative_energy >= energy_threshold)]
+    return f_low, f_high
+
+def detect_bandwidth(signal, fs, energy_threshold=0.99):
+    f, t, Zxx = stft(signal, fs=fs)
+    psd = np.abs(Zxx)**2
+    total_energy = np.sum(psd)
+    cumulative_energy = np.cumsum(np.sum(psd, axis=1)) / total_energy
+    
+    # Exclude DC component (0 Hz)
+    valid_indices = np.where(f > 0)[0]
+    f_low = f[valid_indices][np.argmax(cumulative_energy[valid_indices] > (1 - energy_threshold))]
+    f_high = f[np.argmax(cumulative_energy >= energy_threshold)]
+    return f_low, f_high
+    
+# Step 3: Apply bandpass and lowpass filters
+def bandpass_filter(signal, fs, f_low, f_high):
+    nyquist = 0.5 * fs
+    low = f_low / nyquist
+    high = f_high / nyquist
+    b, a = butter(N=4, Wn=[low, high], btype='band')
+    return filtfilt(b, a, signal)
+
+def lowpass_filter(signal, fs, cutoff):
+    nyquist = 0.5 * fs
+    cutoff_normalized = cutoff / nyquist
+    b, a = butter(N=4, Wn=cutoff_normalized, btype='low')
+    return filtfilt(b, a, signal)
+
+def highpass_filter(signal, fs, cutoff):
+    nyquist = 0.5 * fs
+    cutoff_normalized = cutoff / nyquist
+    b, a = butter(N=4, Wn=cutoff_normalized, btype='high')
+    return filtfilt(b, a, signal)
+
+# Step 4: Replace bandwidth
+def replace_bandwidth(signal1, signal2, fs, f_low, f_high):
+    # Extract effective band from signal1
+    #effective_band = bandpass_filter(signal1, fs, f_low, f_high)
+    effective_band = lowpass_filter(signal1, fs, f_high)
+    # Extract lowpass band from signal2
+    #signal2_lowpass = lowpass_filter(signal2, fs, f_high)
+    signal2_highpass = highpass_filter(signal2, fs, f_high)
+    
+    # Match lengths of the two signals
+    min_length = min(len(effective_band), len(signal2_highpass))
+    effective_band = effective_band[:min_length]
+    signal2_highpass = signal2_highpass[:min_length]
+    
+    # Combine the two signals
+    return signal2_highpass + effective_band
+
+# Step 5: Smooth transitions
+def smooth_transition(signal1, signal2, fs, transition_band=100):
+    fade = np.linspace(0, 1, int(transition_band * fs / 1000))
+    crossfade = np.concatenate([fade, np.ones(len(signal1) - len(fade))])
+    min_length = min(len(signal1), len(signal2))
+    smoothed_signal = (1 - crossfade) * signal2[:min_length] + crossfade * signal1[:min_length]
+    return smoothed_signal
+
+# Step 6: Save audio
+def save_audio(file_path, audio, fs):
+    sf.write(file_path, audio, fs)
+
+
+def bandwidth_sub(low_bandwidth_audio, high_bandwidth_audio, fs=48000):
+    # Detect effective bandwidth of the first signal
+    f_low, f_high = detect_bandwidth(low_bandwidth_audio, fs)
+        
+    # Replace the lower frequency of the second audio
+    substituted_audio = replace_bandwidth(low_bandwidth_audio, high_bandwidth_audio, fs, f_low, f_high)
+
+    # Optional: Smooth the transition
+    smoothed_audio = smooth_transition(substituted_audio, low_bandwidth_audio, fs)
+    return smoothed_audio
+        
+# Main process
+if __name__ == "__main__":
+    low_spectra_dir = 'LJSpeech_22k'
+    upper_spectra_dir = 'LJSpeech_22k_hifi-sr_speech_g_03925000'
+    output_dir = upper_spectra_dir+'_restored'
+    if not os.path.exists(output_dir):
+        os.mkdir(output_dir)
+        
+    filelist = [file for file in os.listdir(low_spectra_dir) if file.endswith('.wav')]
+    for audio_name in filelist:
+        audio1, fs1 = load_audio(low_spectra_dir + "/" + audio_name)  # Source for effective bandwidth
+        audio2, fs2 = load_audio(upper_spectra_dir + "/" + audio_name.replace('.wav', '_generated.wav'))  # Target audio to replace lower frequencies
+
+        if fs1 != 48000 or fs2 != 48000:
+            raise ValueError("Both audio files must have a sampling rate of 48 kHz.")
+
+        # Detect effective bandwidth of the first signal
+        f_low, f_high = detect_bandwidth(audio1, fs1)
+        print(f"Effective bandwidth: {f_low} Hz to {f_high} Hz")
+
+        # Replace the lower frequency of the second audio
+        replaced_audio = replace_bandwidth(audio1, audio2, fs2, f_low, f_high)
+
+        # Optional: Smooth the transition
+        smoothed_audio = smooth_transition(replaced_audio, audio1, fs1)
+
+        # Save the result
+        save_audio(output_dir+"/"+audio_name, smoothed_audio, fs2)

utils/decode.py

@@ -0,0 +1,609 @@
+#!/usr/bin/env python -u
+# -*- coding: utf-8 -*-
+# Authors: Shengkui Zhao, Zexu Pan
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import torch 
+import torch.nn as nn
+import numpy as np
+import os 
+import sys
+import librosa
+import torchaudio
+from utils.misc import power_compress, power_uncompress, stft, istft, compute_fbank
+from utils.bandwidth_sub import bandwidth_sub
+from dataloader.meldataset import mel_spectrogram
+
+# Constant for normalizing audio values
+MAX_WAV_VALUE = 32768.0
+
+def decode_one_audio(model, device, inputs, args):
+    """Decodes audio using the specified model based on the provided network type.
+
+    This function selects the appropriate decoding function based on the specified
+    network in the arguments and processes the input audio data accordingly.
+
+    Args:
+        model (nn.Module): The trained model used for decoding.
+        device (torch.device): The device (CPU or GPU) to perform computations on.
+        inputs (torch.Tensor): Input audio tensor.
+        args (Namespace): Contains arguments for network configuration.
+
+    Returns:
+        list: A list of decoded audio outputs for each speaker.
+    """
+    # Select decoding function based on the network type specified in args
+    if args.network == 'FRCRN_SE_16K':
+        return decode_one_audio_frcrn_se_16k(model, device, inputs, args)
+    elif args.network == 'MossFormer2_SE_48K':
+        return decode_one_audio_mossformer2_se_48k(model, device, inputs, args)
+    elif args.network == 'MossFormerGAN_SE_16K':
+        return decode_one_audio_mossformergan_se_16k(model, device, inputs, args)
+    elif args.network == 'MossFormer2_SS_16K':
+        return decode_one_audio_mossformer2_ss_16k(model, device, inputs, args)
+    elif args.network == 'MossFormer2_SR_48K':
+        return decode_one_audio_mossformer2_sr_48k(model, device, inputs, args)
+    else:
+        print("No network found!")  # Print error message if no valid network is specified
+        return 
+
+def decode_one_audio_mossformer2_ss_16k(model, device, inputs, args):
+    """Decodes audio using the MossFormer2 model for speech separation at 16kHz.
+
+    This function handles the audio decoding process by processing the input tensor
+    in segments, if necessary, and applies the model to obtain separated audio outputs.
+
+    Args:
+        model (nn.Module): The trained MossFormer2 model for decoding.
+        device (torch.device): The device (CPU or GPU) to perform computations on.
+        inputs (torch.Tensor): Input audio tensor of shape (B, T), where B is the batch size
+                              and T is the number of time steps.
+        args (Namespace): Contains arguments for decoding configuration.
+
+    Returns:
+        list: A list of decoded audio outputs for each speaker.
+    """
+    out = []  # Initialize the list to store outputs
+    decode_do_segment = False  # Flag to determine if segmentation is needed
+    window = int(args.sampling_rate * args.decode_window)  # Decoding window length
+    stride = int(window * 0.75)  # Decoding stride if segmentation is used
+    b, t = inputs.shape  # Get batch size and input length
+
+    rms_input = (inputs ** 2).mean() ** 0.5
+
+    # Check if input length exceeds one-time decode length to decide on segmentation
+    if t > args.sampling_rate * args.one_time_decode_length:
+        decode_do_segment = True  # Enable segment decoding for long sequences
+
+    # Pad the inputs to ensure they meet the decoding window length requirements
+    if t < window:
+        inputs = np.concatenate([inputs, np.zeros((inputs.shape[0], window - t))], axis=1)
+    elif t < window + stride:
+        padding = window + stride - t
+        inputs = np.concatenate([inputs, np.zeros((inputs.shape[0], padding))], axis=1)
+    else:
+        if (t - window) % stride != 0:
+            padding = t - (t - window) // stride * stride
+            inputs = np.concatenate([inputs, np.zeros((inputs.shape[0], padding))], axis=1)
+
+    inputs = torch.from_numpy(np.float32(inputs)).to(device)  # Convert inputs to torch tensor and move to device
+    b, t = inputs.shape  # Update batch size and input length after conversion
+
+    # Process the inputs in segments if necessary
+    if decode_do_segment:
+        outputs = np.zeros((args.num_spks, t))  # Initialize output array for each speaker
+        give_up_length = (window - stride) // 2  # Calculate length to give up at each segment
+        current_idx = 0  # Initialize current index for segmentation
+        while current_idx + window <= t:
+            tmp_input = inputs[:, current_idx:current_idx + window]  # Get segment input
+            tmp_out_list = model(tmp_input)  # Forward pass through the model
+            for spk in range(args.num_spks):
+                # Convert output for the current speaker to numpy
+                tmp_out_list[spk] = tmp_out_list[spk][0, :].detach().cpu().numpy()
+                if current_idx == 0:
+                    # For the first segment, use the whole segment minus the give-up length
+                    outputs[spk, current_idx:current_idx + window - give_up_length] = tmp_out_list[spk][:-give_up_length]
+                else:
+                    # For subsequent segments, account for the give-up length at both ends
+                    outputs[spk, current_idx + give_up_length:current_idx + window - give_up_length] = tmp_out_list[spk][give_up_length:-give_up_length]
+            current_idx += stride  # Move to the next segment
+        for spk in range(args.num_spks):
+            out.append(outputs[spk, :])  # Append outputs for each speaker
+    else:
+        # If no segmentation is required, process the entire input
+        out_list = model(inputs)
+        for spk in range(args.num_spks):
+            out.append(out_list[spk][0, :].detach().cpu().numpy())  # Append output for each speaker
+
+    # Normalize the outputs back to the input magnitude for each speaker
+    for spk in range(args.num_spks):
+        rms_out = (out[spk] ** 2).mean() ** 0.5
+        out[spk] = out[spk] / rms_out * rms_input
+    return out  # Return the list of normalized outputs
+
+def decode_one_audio_frcrn_se_16k(model, device, inputs, args):
+    """Decodes audio using the FRCRN model for speech enhancement at 16kHz.
+
+    This function processes the input audio tensor either in segments or as a whole, 
+    depending on the length of the input. The model's inference method is applied 
+    to obtain the enhanced audio output.
+
+    Args:
+        model (nn.Module): The trained FRCRN model used for decoding.
+        device (torch.device): The device (CPU or GPU) to perform computations on.
+        inputs (torch.Tensor): Input audio tensor of shape (B, T), where B is the batch size
+                              and T is the number of time steps.
+        args (Namespace): Contains arguments for decoding configuration.
+
+    Returns:
+        numpy.ndarray: The decoded audio output, which has been enhanced by the model.
+    """
+    decode_do_segment = False  # Flag to determine if segmentation is needed
+
+    window = int(args.sampling_rate * args.decode_window)  # Decoding window length
+    stride = int(window * 0.75)  # Decoding stride for segmenting the input
+    b, t = inputs.shape  # Get batch size (b) and input length (t)
+
+    # Check if input length exceeds one-time decode length to decide on segmentation
+    if t > args.sampling_rate * args.one_time_decode_length:
+        decode_do_segment = True  # Enable segment decoding for long sequences
+
+    # Pad the inputs to meet the decoding window length requirements
+    if t < window:
+        # Pad with zeros if the input length is less than the window size
+        inputs = np.concatenate([inputs, np.zeros((inputs.shape[0], window - t))], axis=1)
+    elif t < window + stride:
+        # Pad the input if its length is less than the window plus stride
+        padding = window + stride - t
+        inputs = np.concatenate([inputs, np.zeros((inputs.shape[0], padding))], axis=1)
+    else:
+        # Ensure the input length is a multiple of the stride
+        if (t - window) % stride != 0:
+            padding = t - (t - window) // stride * stride
+            inputs = np.concatenate([inputs, np.zeros((inputs.shape[0], padding))], axis=1)
+
+    # Convert inputs to a PyTorch tensor and move to the specified device
+    inputs = torch.from_numpy(np.float32(inputs)).to(device)
+    b, t = inputs.shape  # Update batch size and input length after conversion
+
+    # Process the inputs in segments if necessary
+    if decode_do_segment:
+        outputs = np.zeros(t)  # Initialize the output array
+        give_up_length = (window - stride) // 2  # Calculate length to give up at each segment
+        current_idx = 0  # Initialize current index for segmentation
+
+        while current_idx + window <= t:
+            tmp_input = inputs[:, current_idx:current_idx + window]  # Get segment input
+            tmp_output = model.inference(tmp_input).detach().cpu().numpy()  # Inference on segment
+
+            # For the first segment, use the whole segment minus the give-up length
+            if current_idx == 0:
+                outputs[current_idx:current_idx + window - give_up_length] = tmp_output[:-give_up_length]
+            else:
+                # For subsequent segments, account for the give-up length
+                outputs[current_idx + give_up_length:current_idx + window - give_up_length] = tmp_output[give_up_length:-give_up_length]
+
+            current_idx += stride  # Move to the next segment
+    else:
+        # If no segmentation is required, process the entire input
+        outputs = model.inference(inputs).detach().cpu().numpy()  # Inference on full input
+
+    return outputs  # Return the decoded audio output
+
+def decode_one_audio_mossformergan_se_16k(model, device, inputs, args):
+    """Decodes audio using the MossFormerGAN model for speech enhancement at 16kHz.
+
+    This function processes the input audio tensor either in segments or as a whole, 
+    depending on the length of the input. The `_decode_one_audio_mossformergan_se_16k` 
+    function is called to perform the model inference and return the enhanced audio output.
+
+    Args:
+        model (nn.Module): The trained MossFormerGAN model used for decoding.
+        device (torch.device): The device (CPU or GPU) for computation.
+        inputs (torch.Tensor): Input audio tensor of shape (B, T), where B is the batch size 
+                              and T is the number of time steps.
+        args (Namespace): Contains arguments for decoding configuration.
+
+    Returns:
+        numpy.ndarray: The decoded audio output, which has been enhanced by the model.
+    """
+    decode_do_segment = False  # Flag to determine if segmentation is needed
+    window = int(args.sampling_rate * args.decode_window)  # Decoding window length
+    stride = int(window * 0.75)  # Decoding stride for segmenting the input
+    b, t = inputs.shape  # Get batch size (b) and input length (t)
+
+    # Check if input length exceeds one-time decode length to decide on segmentation
+    if t > args.sampling_rate * args.one_time_decode_length:
+        decode_do_segment = True  # Enable segment decoding for long sequences
+
+    # Convert inputs to a PyTorch tensor and move to the specified device
+    inputs = torch.from_numpy(np.float32(inputs)).to(device)
+
+    # Compute normalization factor based on the input
+    norm_factor = torch.sqrt(inputs.size(-1) / torch.sum((inputs ** 2.0), dim=-1))
+
+    b, t = inputs.shape  # Update batch size and input length after conversion
+
+    # Process the inputs in segments if necessary
+    if decode_do_segment:
+        outputs = np.zeros(t)  # Initialize the output array
+        give_up_length = (window - stride) // 2  # Calculate length to give up at each segment
+        current_idx = 0  # Initialize current index for segmentation
+
+        while current_idx + window <= t:
+            tmp_input = inputs[:, current_idx:current_idx + window]  # Get segment input
+            tmp_output = _decode_one_audio_mossformergan_se_16k(model, device, tmp_input, norm_factor, args)  # Inference on segment
+
+            # For the first segment, use the whole segment minus the give-up length
+            if current_idx == 0:
+                outputs[current_idx:current_idx + window - give_up_length] = tmp_output[:-give_up_length]
+            else:
+                # For subsequent segments, account for the give-up length
+                outputs[current_idx + give_up_length:current_idx + window - give_up_length] = tmp_output[give_up_length:-give_up_length]
+
+            current_idx += stride  # Move to the next segment
+
+        return outputs  # Return the accumulated outputs from segments
+    else:
+        # If no segmentation is required, process the entire input
+        return _decode_one_audio_mossformergan_se_16k(model, device, inputs, norm_factor, args)  # Inference on full input
+
+@torch.no_grad()
+def _decode_one_audio_mossformergan_se_16k(model, device, inputs, norm_factor, args):
+    """Processes audio inputs through the MossFormerGAN model for speech enhancement.
+
+    This function performs the following steps:
+    1. Pads the input audio tensor to fit the model requirements.
+    2. Computes a normalization factor for the input tensor.
+    3. Applies Short-Time Fourier Transform (STFT) to convert the audio into the frequency domain.
+    4. Processes the STFT representation through the model to predict the real and imaginary components.
+    5. Uncompresses the predicted spectrogram and applies Inverse STFT (iSTFT) to convert back to time domain audio.
+    6. Normalizes the output audio.
+
+    Args:
+        model (nn.Module): The trained MossFormerGAN model used for decoding.
+        device (torch.device): The device (CPU or GPU) for computation.
+        inputs (torch.Tensor): Input audio tensor of shape (B, T), where B is the batch size and T is the number of time steps.
+        norm_factor (torch.Tensor): A norm tensor to regularize input amplitude
+        args (Namespace): Contains arguments for STFT parameters and normalization.
+
+    Returns:
+        numpy.ndarray: The decoded audio output, which has been enhanced by the model.
+    """
+    input_len = inputs.size(-1)  # Get the length of the input audio
+    nframe = int(np.ceil(input_len / args.win_inc))  # Calculate the number of frames based on window increment
+    padded_len = int(nframe * args.win_inc)  # Calculate the padded length to fit the model
+    padding_len = padded_len - input_len  # Determine how much padding is needed
+
+    # Pad the input audio with the beginning of the input
+    inputs = torch.cat([inputs, inputs[:, :padding_len]], dim=-1)
+
+    # Prepare inputs for STFT by transposing and normalizing
+    inputs = torch.transpose(inputs, 0, 1)  # Change shape for STFT
+    inputs = torch.transpose(inputs * norm_factor, 0, 1)  # Apply normalization factor and transpose back
+
+    # Perform Short-Time Fourier Transform (STFT) on the normalized inputs
+    inputs_spec = stft(inputs, args, center=True, periodic=True, onesided=True)
+    inputs_spec = inputs_spec.to(torch.float32)  # Ensure the spectrogram is in float32 format
+
+    # Compress the power of the spectrogram to improve model performance
+    inputs_spec = power_compress(inputs_spec).permute(0, 1, 3, 2)
+
+    # Pass the compressed spectrogram through the model to get predicted real and imaginary parts
+    out_list = model(inputs_spec)
+    pred_real, pred_imag = out_list[0].permute(0, 1, 3, 2), out_list[1].permute(0, 1, 3, 2)
+
+    # Uncompress the predicted spectrogram to get the magnitude and phase
+    pred_spec_uncompress = power_uncompress(pred_real, pred_imag).squeeze(1)
+
+    # Perform Inverse STFT (iSTFT) to convert back to time domain audio
+    outputs = istft(pred_spec_uncompress, args, center=True, periodic=True, onesided=True)
+
+    # Normalize the output audio by dividing by the normalization factor
+    outputs = outputs.squeeze(0) / norm_factor
+
+    return outputs[:input_len].detach().cpu().numpy()  # Return the output as a numpy array
+
+def decode_one_audio_mossformer2_se_48k(model, device, inputs, args):
+    """Processes audio inputs through the MossFormer2 model for speech enhancement at 48kHz.
+
+    This function decodes audio input using the following steps:
+    1. Normalizes the audio input to a maximum WAV value.
+    2. Checks the length of the input to decide between online decoding and batch processing.
+    3. For longer inputs, processes the audio in segments using a sliding window.
+    4. Computes filter banks and their deltas for the audio segment.
+    5. Passes the filter banks through the model to get a predicted mask.
+    6. Applies the mask to the spectrogram of the audio segment and reconstructs the audio.
+    7. For shorter inputs, processes them in one go without segmentation.
+    
+    Args:
+        model (nn.Module): The trained MossFormer2 model used for decoding.
+        device (torch.device): The device (CPU or GPU) for computation.
+        inputs (torch.Tensor): Input audio tensor of shape (B, T), where B is the batch size and T is the number of time steps.
+        args (Namespace): Contains arguments for sampling rate, window size, and other parameters.
+
+    Returns:
+        numpy.ndarray: The decoded audio output, normalized to the range [-1, 1].
+    """
+    inputs = inputs[0, :]  # Extract the first element from the input tensor
+    input_len = inputs.shape[0]  # Get the length of the input audio
+    inputs = inputs * MAX_WAV_VALUE  # Normalize the input to the maximum WAV value
+
+    # Check if input length exceeds the defined threshold for online decoding
+    if input_len > args.sampling_rate * args.one_time_decode_length:  # 20 seconds
+        online_decoding = True
+        if online_decoding:
+            window = int(args.sampling_rate * args.decode_window)  # Define window length (e.g., 4s for 48kHz)
+            stride = int(window * 0.75)  # Define stride length (e.g., 3s for 48kHz)
+            t = inputs.shape[0]  # Update length after potential padding
+
+            # Pad input if necessary to match window size
+            if t < window:
+                inputs = np.concatenate([inputs, np.zeros(window - t)], 0)
+            elif t < window + stride:
+                padding = window + stride - t
+                inputs = np.concatenate([inputs, np.zeros(padding)], 0)
+            else:
+                if (t - window) % stride != 0:
+                    padding = t - (t - window) // stride * stride
+                    inputs = np.concatenate([inputs, np.zeros(padding)], 0)
+
+            audio = torch.from_numpy(inputs).type(torch.FloatTensor)  # Convert to Torch tensor
+            t = audio.shape[0]  # Update length after conversion
+            outputs = torch.from_numpy(np.zeros(t))  # Initialize output tensor
+            give_up_length = (window - stride) // 2  # Determine length to ignore at the edges
+            dfsmn_memory_length = 0  # Placeholder for potential memory length
+            current_idx = 0  # Initialize current index for sliding window
+
+            # Process audio in sliding window segments
+            while current_idx + window <= t:
+                # Select appropriate segment of audio for processing
+                if current_idx < dfsmn_memory_length:
+                    audio_segment = audio[0:current_idx + window]
+                else:
+                    audio_segment = audio[current_idx - dfsmn_memory_length:current_idx + window]
+
+                # Compute filter banks for the audio segment
+                fbanks = compute_fbank(audio_segment.unsqueeze(0), args)
+                
+                # Compute deltas for filter banks
+                fbank_tr = torch.transpose(fbanks, 0, 1)  # Transpose for delta computation
+                fbank_delta = torchaudio.functional.compute_deltas(fbank_tr)  # First-order delta
+                fbank_delta_delta = torchaudio.functional.compute_deltas(fbank_delta)  # Second-order delta
+                
+                # Transpose back to original shape
+                fbank_delta = torch.transpose(fbank_delta, 0, 1)
+                fbank_delta_delta = torch.transpose(fbank_delta_delta, 0, 1)
+
+                # Concatenate the original filter banks with their deltas
+                fbanks = torch.cat([fbanks, fbank_delta, fbank_delta_delta], dim=1)
+                fbanks = fbanks.unsqueeze(0).to(device)  # Add batch dimension and move to device
+
+                # Pass filter banks through the model
+                Out_List = model(fbanks)
+                pred_mask = Out_List[-1]  # Get the predicted mask from the output
+
+                # Apply STFT to the audio segment
+                spectrum = stft(audio_segment, args)
+                pred_mask = pred_mask.permute(2, 1, 0)  # Permute dimensions for masking
+                masked_spec = spectrum.cpu() * pred_mask.detach().cpu()  # Apply mask to the spectrum
+                masked_spec_complex = masked_spec[:, :, 0] + 1j * masked_spec[:, :, 1]  # Convert to complex form
+
+                # Reconstruct audio from the masked spectrogram
+                output_segment = istft(masked_spec_complex, args, len(audio_segment))
+
+                # Store the output segment in the output tensor
+                if current_idx == 0:
+                    outputs[current_idx:current_idx + window - give_up_length] = output_segment[:-give_up_length]
+                else:
+                    output_segment = output_segment[-window:]  # Get the latest window of output
+                    outputs[current_idx + give_up_length:current_idx + window - give_up_length] = output_segment[give_up_length:-give_up_length]
+                
+                current_idx += stride  # Move to the next segment
+
+    else:
+        # Process the entire audio at once if it is shorter than the threshold
+        audio = torch.from_numpy(inputs).type(torch.FloatTensor)
+        fbanks = compute_fbank(audio.unsqueeze(0), args)
+
+        # Compute deltas for filter banks
+        fbank_tr = torch.transpose(fbanks, 0, 1)
+        fbank_delta = torchaudio.functional.compute_deltas(fbank_tr)
+        fbank_delta_delta = torchaudio.functional.compute_deltas(fbank_delta)
+        fbank_delta = torch.transpose(fbank_delta, 0, 1)
+        fbank_delta_delta = torch.transpose(fbank_delta_delta, 0, 1)
+
+        # Concatenate the original filter banks with their deltas
+        fbanks = torch.cat([fbanks, fbank_delta, fbank_delta_delta], dim=1)
+        fbanks = fbanks.unsqueeze(0).to(device)  # Add batch dimension and move to device
+
+        # Pass filter banks through the model
+        Out_List = model(fbanks)
+        pred_mask = Out_List[-1]  # Get the predicted mask
+        spectrum = stft(audio, args)  # Apply STFT to the audio
+        pred_mask = pred_mask.permute(2, 1, 0)  # Permute dimensions for masking
+        masked_spec = spectrum * pred_mask.detach().cpu()  # Apply mask to the spectrum
+        masked_spec_complex = masked_spec[:, :, 0] + 1j * masked_spec[:, :, 1]  # Convert to complex form
+        
+        # Reconstruct audio from the masked spectrogram
+        outputs = istft(masked_spec_complex, args, len(audio))
+
+    return outputs.numpy() / MAX_WAV_VALUE  # Return the output normalized to [-1, 1]
+
+def get_mel(x, args):
+    """
+    Calls mel_spectrogram() and returns the mel-spectrogram output
+    """
+    
+    return mel_spectrogram(x, args.n_fft, args.num_mels, args.sampling_rate, args.hop_size, args.win_size, args.fmin, args.fmax)
+    
+def decode_one_audio_mossformer2_sr_48k(model, device, inputs, args):
+    """
+    This function decodes a single audio input using a two-stage speech super-resolution model.
+    Supports both offline decoding (for short audio) and online decoding (for long audio)
+    with a sliding window approach.
+
+    Parameters:
+    -----------
+    model : list
+        A list of two-stage models:
+        - model[0]: The transformer-based Mossformer model for feature enhancement.
+        - model[1]: The vocoder for generating high-resolution waveforms.
+    device : str or torch.device
+        The computation device ('cpu' or 'cuda') where the models will run.
+    inputs : torch.Tensor
+        A tensor of shape (batch_size, num_samples) containing low-resolution audio signals.
+        Only the first audio (inputs[0, :]) is processed.
+    args : Namespace
+        An object containing the following attributes:
+        - sampling_rate: Sampling rate of the input audio (e.g., 48,000 Hz).
+        - one_time_decode_length: Maximum duration (in seconds) for offline decoding.
+        - decode_window: Window size (in seconds) for sliding window processing.
+        - Other optional attributes used for Mel spectrogram extraction.
+
+    Returns:
+    --------
+    numpy.ndarray
+        The high-resolution audio waveform as a NumPy array, refined and upsampled.
+    """
+    inputs = inputs[0, :]  # Extract the first element from the input tensor
+    input_len = inputs.shape[0]  # Get the length of the input audio
+    #inputs = inputs * MAX_WAV_VALUE  # Normalize the input to the maximum WAV value
+
+    # Check if input length exceeds the defined threshold for online decoding
+    if input_len > args.sampling_rate * args.one_time_decode_length:  # 20 seconds
+        online_decoding = True
+        if online_decoding:
+            window = int(args.sampling_rate * args.decode_window)  # Define window length (e.g., 4s for 48kHz)
+            stride = int(window * 0.75)  # Define stride length (e.g., 3s for 48kHz)
+            t = inputs.shape[0]  # Update length after potential padding
+
+            # Pad input if necessary to match window size
+            if t < window:
+                inputs = np.concatenate([inputs, np.zeros(window - t)], 0)
+            elif t < window + stride:
+                padding = window + stride - t
+                inputs = np.concatenate([inputs, np.zeros(padding)], 0)
+            else:
+                if (t - window) % stride != 0:
+                    padding = t - (t - window) // stride * stride
+                    inputs = np.concatenate([inputs, np.zeros(padding)], 0)
+
+            audio = torch.from_numpy(inputs).type(torch.FloatTensor)  # Convert to Torch tensor
+            t = audio.shape[0]  # Update length after conversion
+            outputs = torch.from_numpy(np.zeros(t))  # Initialize output tensor
+            give_up_length = (window - stride) // 2  # Determine length to ignore at the edges
+            dfsmn_memory_length = 0  # Placeholder for potential memory length
+            current_idx = 0  # Initialize current index for sliding window
+
+            # Process audio in sliding window segments
+            while current_idx + window <= t:
+                # Select appropriate segment of audio for processing
+                if current_idx < dfsmn_memory_length:
+                    audio_segment = audio[0:current_idx + window]
+                else:
+                    audio_segment = audio[current_idx - dfsmn_memory_length:current_idx + window]
+
+                # Pass filter banks through the model
+                mel_segment = get_mel(audio_segment.unsqueeze(0), args)
+                mossformer_output_segment = model[0](mel_segment.to(device))
+                generator_output_segment = model[1](mossformer_output_segment)
+                generator_output_segment = generator_output_segment.squeeze()
+                offset = len(audio_segment) - len(generator_output_segment)
+                # Store the output segment in the output tensor
+                if current_idx == 0:
+                    outputs[current_idx:current_idx + window - give_up_length] = generator_output_segment[:-give_up_length+offset]
+                else:
+                    generator_output_segment = generator_output_segment[-window:]  # Get the latest window of output
+                    outputs[current_idx + give_up_length:current_idx + window - give_up_length] = generator_output_segment[give_up_length:-give_up_length+offset]
+                
+                current_idx += stride  # Move to the next segment
+
+    else:
+        # Process the entire audio at once if it is shorter than the threshold
+        audio = torch.from_numpy(inputs).type(torch.FloatTensor)
+        mel_input = get_mel(audio.unsqueeze(0), args)
+        mossformer_output = model[0](mel_input.to(device))
+        generator_output = model[1](mossformer_output)
+        outputs = generator_output.squeeze()
+
+    outputs = outputs.cpu().numpy()
+    outputs = bandwidth_sub(inputs, outputs)
+    return outputs
+
+def decode_one_audio_AV_MossFormer2_TSE_16K(model, inputs, args):
+    """Processes video inputs through the AV mossformer2 model with Target speaker extraction (TSE) for decoding at 16kHz.
+
+    This function decodes audio input using the following steps:
+    1. Checks if the input audio length requires segmentation or can be processed in one go.
+    2. If the input audio is long enough, processes it in overlapping segments using a sliding window approach.
+    3. Applies the model to each segment or the entire input, and collects the output.
+
+    Args:
+        model (nn.Module): The trained SpEx model for speech enhancement.
+        inputs (numpy.ndarray): Input audio and visual data
+        args (Namespace): Contains arguments for sampling rate, window size, and other parameters.
+
+    Returns:
+        numpy.ndarray: The decoded audio output as a NumPy array.
+    """
+
+    audio, visual = inputs
+    max_val = np.max(np.abs(audio))
+    if max_val > 1:
+        audio /= max_val
+    
+    b, t = audio.shape  # Get batch size (b) and input length (t)
+
+    decode_do_segement = False  # Flag to determine if segmentation is needed
+    # Check if the input length exceeds the defined threshold for segmentation
+    if t > args.sampling_rate * args.one_time_decode_length:
+        decode_do_segement = True  # Enable segmentation for long inputs
+
+    # Convert inputs to a PyTorch tensor and move to the specified device
+    audio = torch.from_numpy(np.float32(audio)).to(args.device)
+    visual = torch.from_numpy(np.float32(visual)).to(args.device)
+
+    print(audio.shape)
+    print(visual.shape)
+
+    if decode_do_segement:
+        print('********')
+        outputs = np.zeros(t)  # Initialize output array
+        window = args.sampling_rate * args.decode_window  # Window length for processing
+        window_v = 25 * args.decode_window
+        stride = int(window * 0.6)  # Decoding stride for segmenting the input
+        give_up_length = (window - stride) // 2  # Calculate length to give up at each segment
+        current_idx = 0  # Initialize current index for sliding window
+
+        # Process the audio in overlapping segments
+        while current_idx + window < t:
+            tmp_audio = audio[:, current_idx:current_idx + window]  # Select current audio segment
+
+            current_idx_v = int(current_idx/args.sampling_rate*25)  # Select current video segment index
+            tmp_video = visual[:, current_idx_v:current_idx_v + window_v, :, :] # Select current video segment
+            
+            tmp_output = model(tmp_audio, tmp_video).detach().squeeze().cpu().numpy()  # Apply model to the segment
+            
+            # For the first segment, use the whole segment minus the give-up length
+            if current_idx == 0:
+                outputs[current_idx:current_idx + window - give_up_length] = tmp_output[:-give_up_length]
+            else:
+                # For subsequent segments, account for the give-up length
+                outputs[current_idx + give_up_length:current_idx + window - give_up_length] = tmp_output[give_up_length:-give_up_length]
+
+            current_idx += stride  # Move to the next segment
+
+        # Process the last window of audio
+        tmp_audio = audio[:, -window:]
+        tmp_video = visual[:, -window_v:, :, :]
+        tmp_output = model(tmp_audio, tmp_video).detach().squeeze().cpu().numpy()  # Apply model to the segment
+        outputs[-window + give_up_length:] = tmp_output[give_up_length:]
+    else:
+        # Process the entire input at once if segmentation is not needed
+        outputs = model(audio, visual).detach().squeeze().cpu().numpy()
+
+
+    return outputs  # Return the decoded audio output as a NumPy array

utils/misc.py

@@ -0,0 +1,380 @@
+#!/usr/bin/env python -u
+# -*- coding: utf-8 -*-
+
+# Import future compatibility features for Python 2/3
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# Import necessary libraries
+import torch 
+import torch.nn as nn
+import numpy as np
+from joblib import Parallel, delayed
+from pesq import pesq  # PESQ metric for speech quality evaluation
+import os 
+import sys
+import librosa  # Library for audio processing
+import torchaudio  # Library for audio processing with PyTorch
+
+# Constants
+MAX_WAV_VALUE = 32768.0  # Maximum value for WAV files
+EPS = 1e-6  # Small value to avoid division by zero
+
+def read_and_config_file(input_path, decode=0):
+    """Reads input paths from a file or directory and configures them for processing.
+
+    Args:
+        input_path (str): Path to the input directory or file.
+        decode (int): Flag indicating if decoding should occur (1 for decode, 0 for standard read).
+
+    Returns:
+        list: A list of processed paths or dictionaries containing input and label paths.
+    """
+    processed_list = []
+
+    # If decoding is requested, find files in a directory
+    if decode:
+        if os.path.isdir(input_path):
+            processed_list = librosa.util.find_files(input_path, ext="wav")  # Look for WAV files
+            if len(processed_list) == 0:
+                processed_list = librosa.util.find_files(input_path, ext="flac")  # Fallback to FLAC files
+        else:
+            # Read paths from a file
+            with open(input_path) as fid:
+                for line in fid:
+                    path_s = line.strip().split()  # Split line into parts
+                    processed_list.append(path_s[0])  # Append the first part (input path)
+        return processed_list
+
+    # Read input-label pairs from a file
+    with open(input_path) as fid:
+        for line in fid:
+            tmp_paths = line.strip().split()  # Split line into parts
+            if len(tmp_paths) == 3:  # Expecting input, label, and duration
+                sample = {'inputs': tmp_paths[0], 'labels': tmp_paths[1], 'duration': float(tmp_paths[2])}
+            elif len(tmp_paths) == 2:  # Expecting input and label only
+                sample = {'inputs': tmp_paths[0], 'labels': tmp_paths[1]}
+            processed_list.append(sample)  # Append the sample dictionary
+    return processed_list
+
+def load_checkpoint(checkpoint_path, use_cuda):
+    """Loads the model checkpoint from the specified path.
+
+    Args:
+        checkpoint_path (str): Path to the checkpoint file.
+        use_cuda (bool): Flag indicating whether to use CUDA for loading.
+
+    Returns:
+        dict: The loaded checkpoint containing model parameters.
+    """
+    #if use_cuda:
+    #    checkpoint = torch.load(checkpoint_path)  # Load using CUDA
+    #else:
+    checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)  # Load to CPU
+    return checkpoint
+
+def get_learning_rate(optimizer):
+    """Retrieves the current learning rate from the optimizer.
+
+    Args:
+        optimizer (torch.optim.Optimizer): The optimizer instance.
+
+    Returns:
+        float: The current learning rate.
+    """
+    return optimizer.param_groups[0]["lr"]
+
+def reload_for_eval(model, checkpoint_dir, use_cuda):
+    """Reloads a model for evaluation from the specified checkpoint directory.
+
+    Args:
+        model (nn.Module): The model to be reloaded.
+        checkpoint_dir (str): Directory containing checkpoints.
+        use_cuda (bool): Flag indicating whether to use CUDA.
+
+    Returns:
+        None
+    """
+    print('Reloading from: {}'.format(checkpoint_dir))
+    best_name = os.path.join(checkpoint_dir, 'last_best_checkpoint')  # Path to the best checkpoint
+    ckpt_name = os.path.join(checkpoint_dir, 'last_checkpoint')  # Path to the last checkpoint
+    if os.path.isfile(best_name):
+        name = best_name 
+    elif os.path.isfile(ckpt_name):
+        name = ckpt_name
+    else:
+        print('Warning: No existing checkpoint or best_model found!')
+        return
+    
+    with open(name, 'r') as f:
+        model_name = f.readline().strip()  # Read the model name from the checkpoint file
+    checkpoint_path = os.path.join(checkpoint_dir, model_name)  # Construct full checkpoint path
+    print('Checkpoint path: {}'.format(checkpoint_path))
+    checkpoint = load_checkpoint(checkpoint_path, use_cuda)  # Load the checkpoint
+    '''
+    if 'model' in checkpoint:
+        model.load_state_dict(checkpoint['model'], strict=False)  # Load model parameters
+    else:
+        model.load_state_dict(checkpoint, strict=False)
+    '''
+    if 'model' in checkpoint:
+        pretrained_model = checkpoint['model']
+    else:
+        pretrained_model = checkpoint
+    state = model.state_dict()
+    for key in state.keys():
+        if key in pretrained_model and state[key].shape == pretrained_model[key].shape:
+            state[key] = pretrained_model[key]
+        elif key.replace('module.', '') in pretrained_model and state[key].shape == pretrained_model[key.replace('module.', '')].shape:
+             state[key] = pretrained_model[key.replace('module.', '')]
+        elif 'module.'+key in pretrained_model and state[key].shape == pretrained_model['module.'+key].shape:
+             state[key] = pretrained_model['module.'+key]
+        elif self.print: print(f'{key} not loaded')
+    model.load_state_dict(state)
+
+    print('=> Reload well-trained model {} for decoding.'.format(model_name))
+    
+
+def reload_model(model, optimizer, checkpoint_dir, use_cuda=True, strict=True):
+    """Reloads the model and optimizer state from a checkpoint.
+
+    Args:
+        model (nn.Module): The model to be reloaded.
+        optimizer (torch.optim.Optimizer): The optimizer to be reloaded.
+        checkpoint_dir (str): Directory containing checkpoints.
+        use_cuda (bool): Flag indicating whether to use CUDA.
+        strict (bool): If True, requires keys in state_dict to match exactly.
+
+    Returns:
+        tuple: Current epoch and step.
+    """
+    ckpt_name = os.path.join(checkpoint_dir, 'checkpoint')  # Path to the checkpoint file
+    if os.path.isfile(ckpt_name):
+        with open(ckpt_name, 'r') as f:
+            model_name = f.readline().strip()  # Read model name from checkpoint file
+        checkpoint_path = os.path.join(checkpoint_dir, model_name)  # Construct full checkpoint path
+        checkpoint = load_checkpoint(checkpoint_path, use_cuda)  # Load the checkpoint
+        model.load_state_dict(checkpoint['model'], strict=strict)  # Load model parameters
+        optimizer.load_state_dict(checkpoint['optimizer'])  # Load optimizer parameters
+        epoch = checkpoint['epoch']  # Get current epoch
+        step = checkpoint['step']  # Get current step
+        print('=> Reloaded previous model and optimizer.')
+    else:
+        print('[!] Checkpoint directory is empty. Train a new model ...')
+        epoch = 0  # Initialize epoch
+        step = 0  # Initialize step
+    return epoch, step
+
+def save_checkpoint(model, optimizer, epoch, step, checkpoint_dir, mode='checkpoint'):
+    """Saves the model and optimizer state to a checkpoint file.
+
+    Args:
+        model (nn.Module): The model to be saved.
+        optimizer (torch.optim.Optimizer): The optimizer to be saved.
+        epoch (int): Current epoch number.
+        step (int): Current training step number.
+        checkpoint_dir (str): Directory to save the checkpoint.
+        mode (str): Mode of the checkpoint ('checkpoint' or other).
+
+    Returns:
+        None
+    """
+    checkpoint_path = os.path.join(
+        checkpoint_dir, 'model.ckpt-{}-{}.pt'.format(epoch, step))  # Construct checkpoint file path
+    torch.save({'model': model.state_dict(),  # Save model parameters
+                'optimizer': optimizer.state_dict(),  # Save optimizer parameters
+                'epoch': epoch,  # Save epoch
+                'step': step}, checkpoint_path)  # Save checkpoint to file
+
+    # Save the checkpoint name to a file for easy access
+    with open(os.path.join(checkpoint_dir, mode), 'w') as f:
+        f.write('model.ckpt-{}-{}.pt'.format(epoch, step))
+    print("=> Saved checkpoint:", checkpoint_path)
+
+def setup_lr(opt, lr):
+    """Sets the learning rate for all parameter groups in the optimizer.
+
+    Args:
+        opt (torch.optim.Optimizer): The optimizer instance whose learning rate needs to be set.
+        lr (float): The new learning rate to be assigned.
+    
+    Returns:
+        None
+    """
+    for param_group in opt.param_groups:
+        param_group['lr'] = lr  # Update the learning rate for each parameter group
+
+
+def pesq_loss(clean, noisy, sr=16000):
+    """Calculates the PESQ (Perceptual Evaluation of Speech Quality) score between clean and noisy signals.
+
+    Args:
+        clean (ndarray): The clean audio signal.
+        noisy (ndarray): The noisy audio signal.
+        sr (int): Sample rate of the audio signals (default is 16000 Hz).
+
+    Returns:
+        float: The PESQ score or -1 in case of an error.
+    """
+    try:
+        pesq_score = pesq(sr, clean, noisy, 'wb')  # Compute PESQ score
+    except:
+        # PESQ may fail due to silent periods in audio
+        pesq_score = -1  # Assign -1 to indicate error
+    return pesq_score
+
+
+def batch_pesq(clean, noisy):
+    """Computes the PESQ scores for batches of clean and noisy audio signals.
+
+    Args:
+        clean (list of ndarray): List of clean audio signals.
+        noisy (list of ndarray): List of noisy audio signals.
+
+    Returns:
+        torch.FloatTensor: A tensor of normalized PESQ scores or None if any score is -1.
+    """
+    # Parallel processing for calculating PESQ scores for each pair of clean and noisy signals
+    pesq_score = Parallel(n_jobs=-1)(delayed(pesq_loss)(c, n) for c, n in zip(clean, noisy))
+    pesq_score = np.array(pesq_score)  # Convert to NumPy array
+    
+    if -1 in pesq_score:  # Check for errors in PESQ calculations
+        return None
+    
+    # Normalize PESQ scores to a scale of 0 to 1
+    pesq_score = (pesq_score - 1) / 3.5  
+    return torch.FloatTensor(pesq_score).to('cuda')  # Return normalized scores as a tensor
+
+
+def power_compress(x):
+    """Compresses the power of a complex spectrogram.
+
+    Args:
+        x (torch.Tensor): Input tensor with real and imaginary components.
+
+    Returns:
+        torch.Tensor: Compressed magnitude and phase representation of the input.
+    """
+    real = x[..., 0]  # Extract real part
+    imag = x[..., 1]  # Extract imaginary part
+    spec = torch.complex(real, imag)  # Create complex tensor from real and imaginary parts
+    mag = torch.abs(spec)  # Compute magnitude
+    phase = torch.angle(spec)  # Compute phase
+    
+    mag = mag**0.3  # Compress magnitude using power of 0.3
+    real_compress = mag * torch.cos(phase)  # Reconstruct real part
+    imag_compress = mag * torch.sin(phase)  # Reconstruct imaginary part
+    return torch.stack([real_compress, imag_compress], 1)  # Stack compressed parts
+
+
+def power_uncompress(real, imag):
+    """Uncompresses the power of a compressed complex spectrogram.
+
+    Args:
+        real (torch.Tensor): Compressed real component.
+        imag (torch.Tensor): Compressed imaginary component.
+
+    Returns:
+        torch.Tensor: Uncompressed complex spectrogram.
+    """
+    spec = torch.complex(real, imag)  # Create complex tensor from real and imaginary parts
+    mag = torch.abs(spec)  # Compute magnitude
+    phase = torch.angle(spec)  # Compute phase
+    
+    mag = mag**(1./0.3)  # Uncompress magnitude by raising to the power of 1/0.3
+    real_uncompress = mag * torch.cos(phase)  # Reconstruct real part
+    imag_uncompress = mag * torch.sin(phase)  # Reconstruct imaginary part
+    return torch.stack([real_uncompress, imag_uncompress], -1)  # Stack uncompressed parts
+
+
+def stft(x, args, center=False, periodic=False, onesided=None):
+    """Computes the Short-Time Fourier Transform (STFT) of an audio signal.
+
+    Args:
+        x (torch.Tensor): Input audio signal.
+        args (Namespace): Configuration arguments containing window type and lengths.
+        center (bool): Whether to center the window.
+
+    Returns:
+        torch.Tensor: The computed STFT of the input signal.
+    """
+    win_type = args.win_type
+    win_len = args.win_len
+    win_inc = args.win_inc
+    fft_len = args.fft_len
+
+    # Select window type and create window tensor
+    if win_type == 'hamming':
+        window = torch.hamming_window(win_len, periodic=periodic).to(x.device)
+    elif win_type == 'hanning':
+        window = torch.hann_window(win_len, periodic=periodic).to(x.device)
+    else:
+        print(f"In STFT, {win_type} is not supported!")
+        return
+
+    # Compute and return the STFT
+    return torch.stft(x, fft_len, win_inc, win_len, center=center, window=window, onesided=onesided, return_complex=False)
+
+def istft(x, args, slen=None, center=False, normalized=False, periodic=False, onesided=None, return_complex=False):
+    """Computes the inverse Short-Time Fourier Transform (ISTFT) of a complex spectrogram.
+
+    Args:
+        x (torch.Tensor): Input complex spectrogram.
+        args (Namespace): Configuration arguments containing window type and lengths.
+        slen (int, optional): Length of the output signal.
+        center (bool): Whether to center the window.
+        normalized (bool): Whether to normalize the output.
+        onesided (bool, optional): If True, computes only the one-sided transform.
+        return_complex (bool): If True, returns complex output.
+
+    Returns:
+        torch.Tensor: The reconstructed audio signal from the spectrogram.
+    """
+    win_type = args.win_type
+    win_len = args.win_len
+    win_inc = args.win_inc
+    fft_len = args.fft_len
+
+    # Select window type and create window tensor
+    if win_type == 'hamming':
+        window = torch.hamming_window(win_len, periodic=periodic).to(x.device)
+    elif win_type == 'hanning':
+        window = torch.hann_window(win_len, periodic=periodic).to(x.device)
+    else:
+        print(f"In ISTFT, {win_type} is not supported!")
+        return
+
+    try:
+        # Attempt to compute ISTFT
+        output = torch.istft(x, n_fft=fft_len, hop_length=win_inc, win_length=win_len,
+                              window=window, center=center, normalized=normalized,
+                              onesided=onesided, length=slen, return_complex=False)
+    except:
+        # Handle potential errors by converting x to a complex tensor
+        x_complex = torch.view_as_complex(x)
+        output = torch.istft(x_complex, n_fft=fft_len, hop_length=win_inc, win_length=win_len,
+                              window=window, center=center, normalized=normalized,
+                              onesided=onesided, length=slen, return_complex=False)
+    return output
+
+def compute_fbank(audio_in, args):
+    """Computes the filter bank features from an audio signal.
+
+    Args:
+        audio_in (torch.Tensor): Input audio signal.
+        args (Namespace): Configuration arguments containing window length, shift, and sampling rate.
+
+    Returns:
+        torch.Tensor: Computed filter bank features.
+    """
+    frame_length = args.win_len / args.sampling_rate * 1000  # Frame length in milliseconds
+    frame_shift = args.win_inc / args.sampling_rate * 1000  # Frame shift in milliseconds
+
+    # Compute and return filter bank features using Kaldi's implementation
+    return torchaudio.compliance.kaldi.fbank(audio_in, dither=1.0, frame_length=frame_length,
+                                             frame_shift=frame_shift, num_mel_bins=args.num_mels,
+                                             sample_frequency=args.sampling_rate, window_type=args.win_type)
+                                             
+
+

utils/video_process.py

@@ -0,0 +1,361 @@
+import torch
+
+import sys, time, os, tqdm, torch, argparse, glob, subprocess, warnings, cv2, pickle, pdb, math, python_speech_features
+import numpy as np
+from scipy import signal
+from shutil import rmtree
+from scipy.io import wavfile
+from scipy.interpolate import interp1d
+from sklearn.metrics import accuracy_score, f1_score
+import soundfile as sf
+
+from scenedetect.video_manager import VideoManager
+from scenedetect.scene_manager import SceneManager
+from scenedetect.frame_timecode import FrameTimecode
+from scenedetect.stats_manager import StatsManager
+from scenedetect.detectors import ContentDetector
+
+from models.av_mossformer2_tse.faceDetector.s3fd import S3FD
+
+from .decode import decode_one_audio_AV_MossFormer2_TSE_16K
+
+
+
+def process_tse(args, model, device, data_reader, output_wave_dir):
+	video_args = args_param()
+	video_args.model = model
+	video_args.device = device
+	video_args.sampling_rate = args.sampling_rate
+	args.device = device
+	assert args.sampling_rate == 16000
+	with torch.no_grad():
+		for videoPath in data_reader:  # Loop over all video samples
+			savFolder = videoPath.split(os.path.sep)[-1]
+			video_args.savePath = f'{output_wave_dir}/{savFolder.split(".")[0]}/'
+			video_args.videoPath = videoPath
+			main(video_args, args)
+
+
+
+def args_param():
+    warnings.filterwarnings("ignore")
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--nDataLoaderThread',     type=int,   default=10,   help='Number of workers')
+    parser.add_argument('--facedetScale',          type=float, default=0.25, help='Scale factor for face detection, the frames will be scale to 0.25 orig')
+    parser.add_argument('--minTrack',              type=int,   default=50,   help='Number of min frames for each shot')
+    parser.add_argument('--numFailedDet',          type=int,   default=10,   help='Number of missed detections allowed before tracking is stopped')
+    parser.add_argument('--minFaceSize',           type=int,   default=1,    help='Minimum face size in pixels')
+    parser.add_argument('--cropScale',             type=float, default=0.40, help='Scale bounding box')
+    parser.add_argument('--start',                 type=int, default=0,   help='The start time of the video')
+    parser.add_argument('--duration',              type=int, default=0,  help='The duration of the video, when set as 0, will extract the whole video')
+    video_args = parser.parse_args()
+    return video_args
+
+
+# Main function
+def main(video_args, args):
+    # Initialization 
+    video_args.pyaviPath = os.path.join(video_args.savePath, 'py_video')
+    video_args.pyframesPath = os.path.join(video_args.savePath, 'pyframes')
+    video_args.pyworkPath = os.path.join(video_args.savePath, 'pywork')
+    video_args.pycropPath = os.path.join(video_args.savePath, 'py_faceTracks')
+    if os.path.exists(video_args.savePath):
+        rmtree(video_args.savePath)
+    os.makedirs(video_args.pyaviPath, exist_ok = True) # The path for the input video, input audio, output video
+    os.makedirs(video_args.pyframesPath, exist_ok = True) # Save all the video frames
+    os.makedirs(video_args.pyworkPath, exist_ok = True) # Save the results in this process by the pckl method
+    os.makedirs(video_args.pycropPath, exist_ok = True) # Save the detected face clips (audio+video) in this process
+
+    # Extract video
+    video_args.videoFilePath = os.path.join(video_args.pyaviPath, 'video.avi')
+    # If duration did not set, extract the whole video, otherwise extract the video from 'video_args.start' to 'video_args.start + video_args.duration'
+    if video_args.duration == 0:
+        command = ("ffmpeg -y -i %s -qscale:v 2 -threads %d -async 1 -r 25 %s -loglevel panic" % \
+            (video_args.videoPath, video_args.nDataLoaderThread, video_args.videoFilePath))
+    else:
+        command = ("ffmpeg -y -i %s -qscale:v 2 -threads %d -ss %.3f -to %.3f -async 1 -r 25 %s -loglevel panic" % \
+            (video_args.videoPath, video_args.nDataLoaderThread, video_args.start, video_args.start + video_args.duration, video_args.videoFilePath))
+    subprocess.call(command, shell=True, stdout=None)
+    sys.stderr.write(time.strftime("%Y-%m-%d %H:%M:%S") + " Extract the video and save in %s \r\n" %(video_args.videoFilePath))
+
+    # Extract audio
+    video_args.audioFilePath = os.path.join(video_args.pyaviPath, 'audio.wav')
+    command = ("ffmpeg -y -i %s -qscale:a 0 -ac 1 -vn -threads %d -ar 16000 %s -loglevel panic" % \
+        (video_args.videoFilePath, video_args.nDataLoaderThread, video_args.audioFilePath))
+    subprocess.call(command, shell=True, stdout=None)
+    sys.stderr.write(time.strftime("%Y-%m-%d %H:%M:%S") + " Extract the audio and save in %s \r\n" %(video_args.audioFilePath))
+
+    # Extract the video frames
+    command = ("ffmpeg -y -i %s -qscale:v 2 -threads %d -f image2 %s -loglevel panic" % \
+        (video_args.videoFilePath, video_args.nDataLoaderThread, os.path.join(video_args.pyframesPath, '%06d.jpg'))) 
+    subprocess.call(command, shell=True, stdout=None)
+    sys.stderr.write(time.strftime("%Y-%m-%d %H:%M:%S") + " Extract the frames and save in %s \r\n" %(video_args.pyframesPath))
+
+    # Scene detection for the video frames
+    scene = scene_detect(video_args)
+    sys.stderr.write(time.strftime("%Y-%m-%d %H:%M:%S") + " Scene detection and save in %s \r\n" %(video_args.pyworkPath))	
+
+    # Face detection for the video frames
+    faces = inference_video(video_args)
+    sys.stderr.write(time.strftime("%Y-%m-%d %H:%M:%S") + " Face detection and save in %s \r\n" %(video_args.pyworkPath))
+
+    # Face tracking
+    allTracks, vidTracks = [], []
+    for shot in scene:
+        if shot[1].frame_num - shot[0].frame_num >= video_args.minTrack: # Discard the shot frames less than minTrack frames
+            allTracks.extend(track_shot(video_args, faces[shot[0].frame_num:shot[1].frame_num])) # 'frames' to present this tracks' timestep, 'bbox' presents the location of the faces
+    sys.stderr.write(time.strftime("%Y-%m-%d %H:%M:%S") + " Face track and detected %d tracks \r\n" %len(allTracks))
+
+    # Face clips cropping
+    for ii, track in tqdm.tqdm(enumerate(allTracks), total = len(allTracks)):
+        vidTracks.append(crop_video(video_args, track, os.path.join(video_args.pycropPath, '%05d'%ii)))
+    savePath = os.path.join(video_args.pyworkPath, 'tracks.pckl')
+    with open(savePath, 'wb') as fil:
+        pickle.dump(vidTracks, fil)
+    sys.stderr.write(time.strftime("%Y-%m-%d %H:%M:%S") + " Face Crop and saved in %s tracks \r\n" %video_args.pycropPath)
+    fil = open(savePath, 'rb')
+    vidTracks = pickle.load(fil)
+    fil.close()
+
+    # AVSE
+    files = glob.glob("%s/*.avi"%video_args.pycropPath)
+    files.sort()
+
+    est_sources = evaluate_network(files, video_args, args)
+
+    visualization(vidTracks, est_sources, video_args)	
+
+    # combine files in pycrop
+    for idx, file in enumerate(files):
+        print(file)
+        command = f"ffmpeg -i {file} {file[:-9]}orig_{idx}.mp4 ;"
+        command += f"rm {file} ;"
+        command += f"rm {file.replace('.avi', '.wav')} ;"
+
+        command += f"ffmpeg -i {file[:-9]}orig_{idx}.mp4 -i {file[:-9]}est_{idx}.wav -c:v copy -map 0:v:0 -map 1:a:0 -shortest {file[:-9]}est_{idx}.mp4 ;"
+        # command += f"rm {file[:-9]}est_{idx}.wav ;"
+
+        output = subprocess.call(command, shell=True, stdout=None)
+
+    rmtree(video_args.pyworkPath)
+    rmtree(video_args.pyframesPath)
+
+
+
+
+def scene_detect(video_args):
+	# CPU: Scene detection, output is the list of each shot's time duration
+	videoManager = VideoManager([video_args.videoFilePath])
+	statsManager = StatsManager()
+	sceneManager = SceneManager(statsManager)
+	sceneManager.add_detector(ContentDetector())
+	baseTimecode = videoManager.get_base_timecode()
+	videoManager.set_downscale_factor()
+	videoManager.start()
+	sceneManager.detect_scenes(frame_source = videoManager)
+	sceneList = sceneManager.get_scene_list(baseTimecode)
+	savePath = os.path.join(video_args.pyworkPath, 'scene.pckl')
+	if sceneList == []:
+		sceneList = [(videoManager.get_base_timecode(),videoManager.get_current_timecode())]
+	with open(savePath, 'wb') as fil:
+		pickle.dump(sceneList, fil)
+		sys.stderr.write('%s - scenes detected %d\n'%(video_args.videoFilePath, len(sceneList)))
+	return sceneList
+
+def inference_video(video_args):
+	# GPU: Face detection, output is the list contains the face location and score in this frame
+	DET = S3FD(device=video_args.device)
+	flist = glob.glob(os.path.join(video_args.pyframesPath, '*.jpg'))
+	flist.sort()
+	dets = []
+	for fidx, fname in enumerate(flist):
+		image = cv2.imread(fname)
+		imageNumpy = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+		bboxes = DET.detect_faces(imageNumpy, conf_th=0.9, scales=[video_args.facedetScale])
+		dets.append([])
+		for bbox in bboxes:
+		  dets[-1].append({'frame':fidx, 'bbox':(bbox[:-1]).tolist(), 'conf':bbox[-1]}) # dets has the frames info, bbox info, conf info
+		sys.stderr.write('%s-%05d; %d dets\r' % (video_args.videoFilePath, fidx, len(dets[-1])))
+	savePath = os.path.join(video_args.pyworkPath,'faces.pckl')
+	with open(savePath, 'wb') as fil:
+		pickle.dump(dets, fil)
+	return dets
+
+def bb_intersection_over_union(boxA, boxB, evalCol = False):
+	# CPU: IOU Function to calculate overlap between two image
+	xA = max(boxA[0], boxB[0])
+	yA = max(boxA[1], boxB[1])
+	xB = min(boxA[2], boxB[2])
+	yB = min(boxA[3], boxB[3])
+	interArea = max(0, xB - xA) * max(0, yB - yA)
+	boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
+	boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])
+	if evalCol == True:
+		iou = interArea / float(boxAArea)
+	else:
+		iou = interArea / float(boxAArea + boxBArea - interArea)
+	return iou
+
+def track_shot(video_args, sceneFaces):
+	# CPU: Face tracking
+	iouThres  = 0.5     # Minimum IOU between consecutive face detections
+	tracks    = []
+	while True:
+		track     = []
+		for frameFaces in sceneFaces:
+			for face in frameFaces:
+				if track == []:
+					track.append(face)
+					frameFaces.remove(face)
+				elif face['frame'] - track[-1]['frame'] <= video_args.numFailedDet:
+					iou = bb_intersection_over_union(face['bbox'], track[-1]['bbox'])
+					if iou > iouThres:
+						track.append(face)
+						frameFaces.remove(face)
+						continue
+				else:
+					break
+		if track == []:
+			break
+		elif len(track) > video_args.minTrack:
+			frameNum    = np.array([ f['frame'] for f in track ])
+			bboxes      = np.array([np.array(f['bbox']) for f in track])
+			frameI      = np.arange(frameNum[0],frameNum[-1]+1)
+			bboxesI    = []
+			for ij in range(0,4):
+				interpfn  = interp1d(frameNum, bboxes[:,ij])
+				bboxesI.append(interpfn(frameI))
+			bboxesI  = np.stack(bboxesI, axis=1)
+			if max(np.mean(bboxesI[:,2]-bboxesI[:,0]), np.mean(bboxesI[:,3]-bboxesI[:,1])) > video_args.minFaceSize:
+				tracks.append({'frame':frameI,'bbox':bboxesI})
+	return tracks
+
+def crop_video(video_args, track, cropFile):
+	# CPU: crop the face clips
+	flist = glob.glob(os.path.join(video_args.pyframesPath, '*.jpg')) # Read the frames
+	flist.sort()
+	vOut = cv2.VideoWriter(cropFile + 't.avi', cv2.VideoWriter_fourcc(*'XVID'), 25, (224,224))# Write video
+	dets = {'x':[], 'y':[], 's':[]}
+	for det in track['bbox']: # Read the tracks
+		dets['s'].append(max((det[3]-det[1]), (det[2]-det[0]))/2) 
+		dets['y'].append((det[1]+det[3])/2) # crop center x 
+		dets['x'].append((det[0]+det[2])/2) # crop center y
+	dets['s'] = signal.medfilt(dets['s'], kernel_size=13)  # Smooth detections 
+	dets['x'] = signal.medfilt(dets['x'], kernel_size=13)
+	dets['y'] = signal.medfilt(dets['y'], kernel_size=13)
+	for fidx, frame in enumerate(track['frame']):
+		cs  = video_args.cropScale
+		bs  = dets['s'][fidx]   # Detection box size
+		bsi = int(bs * (1 + 2 * cs))  # Pad videos by this amount 
+		image = cv2.imread(flist[frame])
+		frame = np.pad(image, ((bsi,bsi), (bsi,bsi), (0, 0)), 'constant', constant_values=(110, 110))
+		my  = dets['y'][fidx] + bsi  # BBox center Y
+		mx  = dets['x'][fidx] + bsi  # BBox center X
+		face = frame[int(my-bs):int(my+bs*(1+2*cs)),int(mx-bs*(1+cs)):int(mx+bs*(1+cs))]
+		vOut.write(cv2.resize(face, (224, 224)))
+	audioTmp    = cropFile + '.wav'
+	audioStart  = (track['frame'][0]) / 25
+	audioEnd    = (track['frame'][-1]+1) / 25
+	vOut.release()
+	command = ("ffmpeg -y -i %s -async 1 -ac 1 -vn -acodec pcm_s16le -ar 16000 -threads %d -ss %.3f -to %.3f %s -loglevel panic" % \
+		      (video_args.audioFilePath, video_args.nDataLoaderThread, audioStart, audioEnd, audioTmp)) 
+	output = subprocess.call(command, shell=True, stdout=None) # Crop audio file
+	_, audio = wavfile.read(audioTmp)
+	command = ("ffmpeg -y -i %st.avi -i %s -threads %d -c:v copy -c:a copy %s.avi -loglevel panic" % \
+			  (cropFile, audioTmp, video_args.nDataLoaderThread, cropFile)) # Combine audio and video file
+	output = subprocess.call(command, shell=True, stdout=None)
+	os.remove(cropFile + 't.avi')
+	return {'track':track, 'proc_track':dets}
+
+
+def evaluate_network(files, video_args, args):
+
+	est_sources = []
+	for file in tqdm.tqdm(files, total = len(files)):
+
+		fileName = os.path.splitext(file.split(os.path.sep)[-1])[0] # Load audio and video
+		audio, _ = sf.read(os.path.join(video_args.pycropPath, fileName + '.wav'), dtype='float32')
+
+		video = cv2.VideoCapture(os.path.join(video_args.pycropPath, fileName + '.avi'))
+		videoFeature = []
+		while video.isOpened():
+			ret, frames = video.read()
+			if ret == True:
+				face = cv2.cvtColor(frames, cv2.COLOR_BGR2GRAY)
+				face = cv2.resize(face, (224,224))
+				face = face[int(112-(112/2)):int(112+(112/2)), int(112-(112/2)):int(112+(112/2))]
+				videoFeature.append(face)
+			else:
+				break
+
+		video.release()
+		visual = np.array(videoFeature)/255.0
+		visual = (visual - 0.4161)/0.1688
+
+		length = int(audio.shape[0]/16000*25)
+		if visual.shape[0] < length:
+			visual = np.pad(visual, ((0,int(length - visual.shape[0])),(0,0),(0,0)), mode = 'edge')
+
+		audio /= np.max(np.abs(audio))
+		audio = np.expand_dims(audio, axis=0)
+		visual = np.expand_dims(visual, axis=0)
+
+		inputs = (audio, visual)
+		est_source = decode_one_audio_AV_MossFormer2_TSE_16K(video_args.model, inputs, args)
+
+		est_sources.append(est_source)
+
+	return est_sources
+
+def visualization(tracks, est_sources, video_args):
+	# CPU: visulize the result for video format
+	flist = glob.glob(os.path.join(video_args.pyframesPath, '*.jpg'))
+	flist.sort()
+	
+
+	for idx, audio in enumerate(est_sources):
+		max_value = np.max(np.abs(audio))
+		if max_value >1:
+			audio /= max_value
+		sf.write(video_args.pycropPath +'/est_%s.wav' %idx, audio, 16000)
+
+	for tidx, track in enumerate(tracks):
+		faces = [[] for i in range(len(flist))]
+		for fidx, frame in enumerate(track['track']['frame'].tolist()):
+			faces[frame].append({'track':tidx, 's':track['proc_track']['s'][fidx], 'x':track['proc_track']['x'][fidx], 'y':track['proc_track']['y'][fidx]})
+	
+		firstImage = cv2.imread(flist[0])
+		fw = firstImage.shape[1]
+		fh = firstImage.shape[0]
+		vOut = cv2.VideoWriter(os.path.join(video_args.pyaviPath, 'video_only.avi'), cv2.VideoWriter_fourcc(*'XVID'), 25, (fw,fh))
+		for fidx, fname in tqdm.tqdm(enumerate(flist), total = len(flist)):
+			image = cv2.imread(fname)
+			for face in faces[fidx]:
+				cv2.rectangle(image, (int(face['x']-face['s']), int(face['y']-face['s'])), (int(face['x']+face['s']), int(face['y']+face['s'])),(0,255,0),10)
+			vOut.write(image)
+		vOut.release()
+
+		command = ("ffmpeg -y -i %s -i %s -threads %d -c:v copy -c:a copy %s -loglevel panic" % \
+			(os.path.join(video_args.pyaviPath, 'video_only.avi'), (video_args.pycropPath +'/est_%s.wav' %tidx), \
+			video_args.nDataLoaderThread, os.path.join(video_args.pyaviPath,'video_out_%s.avi'%tidx))) 
+		output = subprocess.call(command, shell=True, stdout=None)
+
+
+
+
+		command = "ffmpeg -i %s %s ;" % (
+					os.path.join(video_args.pyaviPath, 'video_out_%s.avi' % tidx),
+					os.path.join(video_args.pyaviPath, 'video_est_%s.mp4' % tidx)
+				)
+		command += f"rm {os.path.join(video_args.pyaviPath, 'video_out_%s.avi' % tidx)}"
+		output = subprocess.call(command, shell=True, stdout=None)
+
+
+	command = "ffmpeg -i %s %s ;" % (
+				os.path.join(video_args.pyaviPath, 'video.avi'),
+				os.path.join(video_args.pyaviPath, 'video_orig.mp4')
+			)
+	command += f"rm {os.path.join(video_args.pyaviPath, 'video_only.avi')} ;"
+	command += f"rm {os.path.join(video_args.pyaviPath, 'video.avi')} ;"
+	command += f"rm {os.path.join(video_args.pyaviPath, 'audio.wav')} ;"
+	output = subprocess.call(command, shell=True, stdout=None)