Delete utils

utils/decode.py

@@ -1,532 +0,0 @@
-#!/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
-
-# 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)
-    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 to the maximum absolute value for each speaker
-    '''
-    max_abs = 0
-    for spk in range(args.num_spks):
-        if max_abs < max(abs(out[spk])):
-            max_abs = max(abs(out[spk]))
-    for spk in range(args.num_spks):
-        out[spk] = out[spk] / max_abs  # Normalize output by max absolute value
-    '''
-    # 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
-
-    #normalize outputs
-    #max_abs = max(max(abs(outputs)), 1e-6)
-    #outputs = outputs / max_abs
-    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
-
-    # 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 = _decode_one_audio_mossformergan_se_16k(model, device, tmp_input, 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, args)  # Inference on full input
-
-def _decode_one_audio_mossformergan_se_16k(model, device, inputs, 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.
-        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)
-
-    # Compute normalization factor based on the input
-    c = torch.sqrt(inputs.size(-1) / torch.sum((inputs ** 2.0), dim=-1))
-
-    # Prepare inputs for STFT by transposing and normalizing
-    inputs = torch.transpose(inputs, 0, 1)  # Change shape for STFT
-    inputs = torch.transpose(inputs * c, 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)
-    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)
-
-    # Normalize the output audio by dividing by the normalization factor
-    outputs = outputs.squeeze(0) / c
-
-    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))
-
-    outpus = outputs.numpy() / MAX_WAV_VALUE  # Return the output normalized to [-1, 1]
-    #normalize outputs
-    max_abs = max(max(abs(outputs)), 1e-6) 
-    outputs = outputs / max_abs
-    
-    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

@@ -1,378 +0,0 @@
-#!/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 = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
-    #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]
-    model.load_state_dict(state)
-    print('=> Reloaded 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):
-    """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=False).to(x.device)
-    elif win_type == 'hanning':
-        window = torch.hann_window(win_len, periodic=False).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, return_complex=False)
-
-
-def istft(x, args, slen=None, center=False, normalized=False, onsided=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.
-        onsided (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=False).to(x.device)
-    elif win_type == 'hanning':
-        window = torch.hann_window(win_len, periodic=False).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=onsided, 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=onsided, 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

@@ -1,363 +0,0 @@
-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('/')[-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('/')[-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.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)