Audio Data Conversion and Preprocessing Pipeline

validated_clips_folder = r'Datasets\Data\cv-corpus-19.0-2024-09-13\ha\validated_clips'

# Ensure the validated_clips folder exists
os.makedirs(validated_clips_folder, exist_ok=True)

try:
    # Load the JSON data
    with open(json_file_path, 'r', encoding='utf-8') as f:
        validated_json_data = json.load(f)

    # Extract the list of files from the JSON
    files_in_json = {entry['FILE_PATH'] for entry in validated_json_data}

    # Copy files from clips folder to validated_clips folder
    for file_name in files_in_json:
        source_file = os.path.join(clips_folder, file_name)
        destination_file = os.path.join(validated_clips_folder, file_name)
        
        # Check if the source file exists (in case of discrepancies)
        if os.path.exists(source_file):
            shutil.copy2(source_file, destination_file)
        else:
            print(f"File not found: {source_file}")

    print(f"Clips from the JSON file have been copied to: {validated_clips_folder}")

except FileNotFoundError as e:
    print(f"Error: {e}")

print(f"Validated clips folder created at: {os.path.abspath(validated_clips_folder)}")

# PATH OF VALIDATED.tsv
validated_tsv = 'validated.tsv'

# Code for Language
language_code = 'uz'

def read_tsv_with_missing_columns(tsv_path):
    """
    Reads a .tsv file line by line and handles missing columns by filling them with empty strings.
    """
    rows = []
    with open(tsv_path, 'r', encoding='utf-8') as file:
        # Read the header
        header = file.readline().strip().split('\t')
        total_lines = 0  # Counter for total lines (excluding header)
        for line in file:
            total_lines += 1  # Increment line counter
            # Split the line by tabs and handle missing columns
            row = line.strip().split('\t')
            if len(row) < len(header):
                # Fill missing columns with empty strings
                row.extend([''] * (len(header) - len(row)))
            rows.append(row)
    
    # Convert rows into a DataFrame
    df = pd.DataFrame(rows, columns=header)
    return df, total_lines

def get_audio_duration(file_path):
    ext = os.path.splitext(file_path)[1].lower()
    try:
        if ext == '.mp3':
            audio = MP3(file_path)
        else:
            return None
        return audio.info.length
    except Exception as e:
        print(f"Error processing {file_path}: {e}")
        return None

def process_audio_file(file_path, min_duration=30):
    duration = get_audio_duration(file_path)
    if duration and duration > min_duration:
        print(f"File: {file_path}, Duration: {duration:.2f} seconds")
        return file_path
    return None

def get_validated_audio_files(tsv_path, folder_path):
    """
    Reads the .tsv file using `read_tsv_with_missing_columns` and filters for .mp3 files.
    """
    df, total_lines = read_tsv_with_missing_columns(tsv_path)
    
    # Ensure the 'path' column exists
    if 'path' not in df.columns:
        raise ValueError(f"The TSV file must contain a 'path' column.")
    
    # Filter for .mp3 files and construct full paths
    audio_files = df['path'].apply(lambda x: os.path.join(folder_path, x)).tolist()
    audio_files = [file for file in audio_files if file.endswith('.mp3')]

    print(f"Found {len(audio_files)} validated .mp3 files in the TSV (Total lines in file: {total_lines}).")
    return audio_files

def find_long_audio_files(audio_files, min_duration=30):
    global long_audio_found 

    print(f"Found {len(audio_files)} audio files to process.")

    long_files = []
    with ThreadPoolExecutor(max_workers=max(1, os.cpu_count() // 2)) as executor:
        futures = {executor.submit(process_audio_file, file_path, min_duration): file_path for file_path in audio_files}
        for future in as_completed(futures):
            result = future.result()
            if result:
                long_files.append(result)

    print(f"Found {len(long_files)} audio files longer than {min_duration} seconds.")

    if long_files:
        with open(LONG_AUDIO_FILE, "w") as f:
            f.write("\n".join(long_files) + "\n")
        print(f"List of long audio files saved to {LONG_AUDIO_FILE}")
        long_audio_found = True 
        print(f"long_audio_found set to TRUE.")
    else:
        print("No audio files longer than 30 seconds found.")

# Running the code
if __name__ == "__main__":
    tsv_path = "validated.tsv"  # Path to the .tsv file
    folder_path = "clips"       # Folder where audio files are stored

    # Get the list of validated .mp3 files
    validated_audio_files = get_validated_audio_files(tsv_path, folder_path)

    # Process only the validated files
    find_long_audio_files(validated_audio_files)

# Function to read and process the .tsv file while preserving all rows
def read_tsv_with_missing_columns(tsv_path):
    # Read the file line by line to handle malformed rows
    rows = []
    with open(tsv_path, 'r', encoding='utf-8') as file:
        # Read the header
        header = file.readline().strip().split('\t')
        total_lines = 0  # Counter for total lines (excluding header)
        for line in file:
            total_lines += 1  # Increment line counter
            # Split the line by tabs and handle missing columns
            row = line.strip().split('\t')
            if len(row) < len(header):
                # Fill missing columns with empty strings
                row.extend([''] * (len(header) - len(row)))
            rows.append(row)
    
    # Convert rows into a DataFrame
    df = pd.DataFrame(rows, columns=header)
    return df, total_lines

# Function to convert audio files to 16kHz WAV format
def convert_to_wav(input_file, output_file):
    try:
        # Load the audio file and convert it to 16kHz mono
        audio = AudioSegment.from_file(input_file)
        audio = audio.set_frame_rate(16000).set_channels(1)
        # Export the audio to WAV format
        audio.export(output_file, format="wav")
        return f"Finished converting and moving: {input_file}"
    except Exception as e:
        return f"Error converting {input_file}: {e}"

# Main processing function (INCREASE MAX_WORKERS FOR FASTER CONVERSION)
def process_tsv(tsv_path, audio_folder, output_folder, max_workers =  max(1, os.cpu_count() // 2)):
    # Read the .tsv file with the new logic and get the total line count
    data, total_lines = read_tsv_with_missing_columns(tsv_path)

    # Ensure the output folder exists
    os.makedirs(output_folder, exist_ok=True)

    # Create a ThreadPoolExecutor to parallelize the conversion tasks
    with ThreadPoolExecutor(max_workers= max_workers) as executor:
        # Submit tasks to the executor
        futures = []
        for index, row in data.iterrows():
            input_file = os.path.join(audio_folder, row['path'])
            output_file = os.path.join(output_folder, f"{os.path.splitext(row['path'])[0]}.wav")

            if os.path.exists(input_file):
                # Submit the conversion task to the executor
                future = executor.submit(convert_to_wav, input_file, output_file)
                futures.append(future)
            else:
                print(f"File not found: {input_file}")

        # Process the results as they complete
        for future in as_completed(futures):
            try:
                result = future.result()
                print(result)
            except Exception as e:
                print(f"Error during conversion: {e}")

    # Validate the number of rows processed
    print(f"Total rows in DataFrame: {len(data)}")
    print(f"Total lines in TSV file (excluding header): {total_lines}")

    if len(data) == total_lines:
        print("Validation successful: The number of rows matches the number of lines in the TSV file.")
    else:
        print("Validation failed: The number of rows does not match the number of lines in the TSV file.")



if __name__ == "__main__":
    # Define paths
    tsv_path = validated_tsv  # Path to your .tsv file
    audio_folder = "clips"  # Path to the folder containing audio files
    output_folder = f"validated_{language_code}_wav"  # Output folder for converted .wav files

    # Process the .tsv and convert files
    process_tsv(tsv_path, audio_folder, output_folder)

# Function to process the TSV and create splits
def process_tsv_for_splits(tsv_path, output_folder):
    # Read the .tsv file with the row-preserving logic
    data, total_lines = read_tsv_with_missing_columns(tsv_path)

    # Ensure the required columns exist
    if 'path' not in data.columns or 'sentence' not in data.columns:
        raise ValueError("The TSV file must contain 'path' and 'sentence' columns.")

    # Extract only the 'path' and 'sentence' columns
    data = data[['path', 'sentence']]

    # Rename 'sentence' to 'text'
    data = data.rename(columns={'sentence': 'text'})

    # Modify the 'path' column: replace .mp3 with .wav, prepend the output folder, and use forward slashes
    data['path'] = data['path'].apply(
        lambda x: os.path.join(output_folder, os.path.splitext(x)[0] + '.wav').replace('\\', '/')
    )

    # Check if the flag is True and read the .txt file if it exists
    if long_audio_found:
        excluded_files = read_excluded_files(LONG_AUDIO_FILE)
        if excluded_files:
            print("Excluding files longer than 30 seconds from the TSV processing.")
            data = data[~data['path'].isin(excluded_files)]  # Exclude files in the excluded_files list
        else:
            print("No files to exclude: The .txt file is empty or does not exist.")
    else:
        print("No long audio files found. Proceeding without exclusions.")

    # Add a new 'split' column with 80/10/10 train/dev/test split
    np.random.seed(42)  # Set a seed for reproducibility
    split_proportions = [0.8, 0.1, 0.1]  # Train/Dev/Test proportions
    splits = np.random.choice(['train', 'dev', 'test'], size=len(data), p=split_proportions)
    data['split'] = splits

    # Export the modified DataFrame to a new .tsv file in the current directory
    output_tsv_path = f'{language_code}_validated_audio_splits.tsv'  # Save in the current directory
    data.to_csv(output_tsv_path, sep='\t', index=False)
    print(f"Exported new TSV file to: {os.path.abspath(output_tsv_path)}")

    # Validate the number of rows processed
    print(f"Total rows in DataFrame: {len(data)}")
    print(f"Total lines in TSV file (excluding header): {total_lines}")

    if len(data) == total_lines:
        print("Validation successful: The number of rows matches the number of lines in the TSV file.")
    else:
        print("Validation failed: The number of rows does not match the number of lines in the TSV file.")

def read_excluded_files(file_path):
    if not os.path.exists(file_path):
        print(f"File {file_path} does not exist.")
        return []
    
    with open(file_path, 'r', encoding='utf-8') as file:
        excluded_files = file.read().splitlines()
    return excluded_files


# Check paths here - rename output folder
if __name__ == "__main__":
    # Define paths
    tsv_path = validated_tsv  # Path to your .tsv file
    output_folder = f"validated_{language_code}_wavs"  # Folder where the .wav files are stored

    # Process the .tsv and export the new file
    process_tsv_for_splits(tsv_path, output_folder)

# Function to create packets of 30k rows with 80/10/10 split
def create_packets(tsv_path, packet_size=30000, train_ratio=0.8, dev_ratio=0.1, test_ratio=0.1):
    # Read the .tsv file with the row-preserving logic
    data, total_lines = read_tsv_with_missing_columns(tsv_path)

    # Ensure the required columns exist
    if 'path' not in data.columns or 'text' not in data.columns:
        raise ValueError("The TSV file must contain 'path' and 'text' columns.")

    # Shuffle the data to ensure randomness
    data = data.sample(frac=1, random_state=42).reset_index(drop=True)

    # Calculate the number of packets
    num_packets = (len(data) // packet_size) + (1 if len(data) % packet_size != 0 else 0)

    # Create packets
    for i in range(num_packets):
        # Get the current packet
        start_idx = i * packet_size
        end_idx = start_idx + packet_size
        packet = data[start_idx:end_idx]

        # Add a new 'split' column with 80/10/10 train/dev/test split
        np.random.seed(42 + i)  # Set a seed for reproducibility (unique for each packet)
        split_proportions = [train_ratio, dev_ratio, test_ratio]
        splits = np.random.choice(['train', 'dev', 'test'], size=len(packet), p=split_proportions)
        packet['split'] = splits

        # Export the packet to a new .tsv file
        output_tsv_path = f'{language_code}_packet_{i + 1}.tsv'
        packet.to_csv(output_tsv_path, sep='\t', index=False)
        print(f"Exported packet {i + 1} to: {os.path.abspath(output_tsv_path)}")

# UNCOMMENT BELOW TO USE - CHECK TSV PATH TO MATCH NEW .TSV

if __name__ == "__main__":
    # Define paths
    tsv_path = f"{language_code}_validated_audio_splits.tsv"  # Path to your .tsv file

    # Create packets of 30k rows each with 80/10/10 split
    create_packets(tsv_path)

def copy_audio_files(tsv_path):
    # Load the TSV file
    df = pd.read_csv(tsv_path, sep='\t')

    # Ensure 'path' column exists
    if 'path' not in df.columns:
        raise ValueError("The TSV file does not contain a 'path' column.")

    # Define destination folder
    base_name = os.path.splitext(os.path.basename(tsv_path))[0]
    dest_folder = f"{base_name}_wavs"

    # Create destination folder if it doesn't exist
    os.makedirs(dest_folder, exist_ok=True)

    # Function to copy a single file
    def copy_file(audio_path):
        try:
            shutil.copy(audio_path, dest_folder)
            return audio_path, True
        except Exception as e:
            print(f"Warning: {audio_path} not found or could not be copied. Error: {e}")
            return audio_path, False

    # Use ThreadPoolExecutor to copy files in parallel
    with ThreadPoolExecutor(max_workers=os.cpu_count() * 2) as executor:
        future_to_file = {executor.submit(copy_file, audio_path): audio_path for audio_path in df['path']}
        
        success_count = 0
        for future in as_completed(future_to_file):
            audio_path, success = future.result()
            if success:
                success_count += 1

    print(f"Copied {success_count} files to {dest_folder}")

# Example usage
tsv_file_path = "uz_packet_1.tsv"  # Replace with your actual TSV file path
copy_audio_files(tsv_file_path)