igna/LuxTTS
1
0
Fork 0
Code Activity
Files
c6efaff7f4759a4a9b33f435cea367c1944aea47
LuxTTS/zipvoice/bin/prepare_dataset.py
Yatharth Sharma b51a0820b3 Initial commit
2026-01-23 15:12:06 -05:00

275 lines
8.7 KiB
Python
Raw Blame History

#!/usr/bin/env python3
# Copyright 2025 Xiaomi Corp. (authors: Han Zhu)
#
# See ../../../../LICENSE for clarification regarding multiple authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This script generates lhotse manifest files from TSV files for custom datasets.
Each line of the TSV files should be in one of the following formats:
1. "{uniq_id}\t{text}\t{wav_path}" if the text corresponds to the full wav",
2. "{uniq_id}\t{text}\t{wav_path}\t{start_time}\t{end_time} if text corresponds
to part of the wav. The start_time and end_time specify the start and end
times of the text within the wav, which should be in seconds.
Note: {uniq_id} must be unique for each line.
Usage:
Suppose you have two TSV files: "custom_train.tsv" and "custom_dev.tsv",
where "custom" is your dataset name, "train"/"dev" are used for training and
validation respectively.
(1) Prepare the training data
python3 -m zipvoice.bin.prepare_dataset \
--tsv-path data/raw/custom_train.tsv \
--prefix "custom" \
--subset "train" \
--num-jobs 20 \
--output-dir "data/manifests"
The output file would be "data/manifests/custom_cuts_train.jsonl.gz".
(2) Prepare the validation data
python3 -m zipvoice.bin.prepare_dataset \
--tsv-path data/raw/custom_dev.tsv \
--prefix "custom" \
--subset "dev" \
--num-jobs 1 \
--output-dir "data/manifests"
The output file would be "data/manifests/custom_cuts_dev.jsonl.gz".
"""
import argparse
import logging
import re
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import List, Optional, Tuple
from lhotse import CutSet, validate_recordings_and_supervisions
from lhotse.audio import Recording, RecordingSet
from lhotse.qa import fix_manifests
from lhotse.supervision import SupervisionSegment, SupervisionSet
from lhotse.utils import Pathlike
from tqdm.auto import tqdm
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--tsv-path",
type=str,
help="The path of the tsv file. Each line should be in the format: "
"{uniq_id}\t{text}\t{wav_path}\t{start_time}\t{end_time} "
"if text corresponds to part of the wav or {uniq_id}\t{text}\t{wav_path} "
"if the text corresponds to the full wav",
)
parser.add_argument(
"--prefix",
type=str,
default="custom",
help="Prefix of the output manifest file.",
)
parser.add_argument(
"--subset",
type=str,
default="train",
help="Subset name manifest file, typically train or dev.",
)
parser.add_argument(
"--num-jobs",
type=int,
default=20,
help="Number of jobs to processing.",
)
parser.add_argument(
"--output-dir",
type=str,
default="data/manifests",
help="The destination directory of manifest files.",
)
parser.add_argument(
"--sampling-rate",
type=int,
default=24000,
help="The target sampling rate.",
)
return parser.parse_args()
def _parse_recording(
wav_path: str,
) -> Tuple[Recording, str]:
"""
:param wav_path: Path to the audio file
:return: a tuple of "recording" and "recording_id"
"""
recording_id = wav_path.replace("/", "_").replace(".", "_")
recording = Recording.from_file(path=wav_path, recording_id=recording_id)
return recording, recording_id
def _parse_supervision(
supervision: List, recording_dict: dict
) -> Optional[SupervisionSegment]:
"""
:param line: A line from the TSV file
:param recording_dict: Dictionary mapping recording IDs to Recording objects
:return: A SupervisionSegment object
"""
uniq_id, text, wav_path, start, end = supervision
try:
recording_id = wav_path.replace("/", "_").replace(".", "_")
recording = recording_dict[recording_id]
duration = end - start if end is not None else recording.duration
assert duration <= recording.duration, f"Duration {duration} is greater than "
f"recording duration {recording.duration}"
text = re.sub("_", " ", text) # "_" is treated as padding symbol
text = re.sub(r"\s+", " ", text) # remove extra whitespace
return SupervisionSegment(
id=f"{uniq_id}",
recording_id=recording.id,
start=start,
duration=duration,
channel=recording.channel_ids,
text=text.strip(),
)
except Exception as e:
logging.warning(f"Error processing line: {e}")
return None
def prepare_dataset(
tsv_path: Pathlike,
prefix: str,
subset: str,
sampling_rate: int,
num_jobs: int,
output_dir: Pathlike,
):
"""
Returns the manifests which consist of the Recordings and Supervisions
:param tsv_path: Path to the TSV file
:param output_dir: Path where to write the manifests
:param num_jobs: Number of processes for parallel processing
:return: The CutSet containing the data
"""
logging.info(f"Preparing {prefix} dataset {subset} subset.")
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
file_name = f"{prefix}_cuts_{subset}.jsonl.gz"
if (output_dir / file_name).is_file():
logging.info(f"{file_name} exists, skipping.")
return
# Step 1: Read all unique recording paths
recordings_path_set = set()
supervision_list = list()
with open(tsv_path, "r") as fr:
for line in fr:
items = line.strip().split("\t")
if len(items) == 3:
uniq_id, text, wav_path = items
start, end = 0, None
elif len(items) == 5:
uniq_id, text, wav_path, start, end = items
start, end = float(start), float(end)
else:
raise ValueError(
f"Invalid line format: {line},"
"requries to be 3 columns or 5 columns"
)
recordings_path_set.add(wav_path)
supervision_list.append((uniq_id, text, wav_path, start, end))
logging.info("Starting to process recordings...")
# Step 2: Process recordings
futures = []
recording_dict = {}
with ThreadPoolExecutor(max_workers=num_jobs) as ex:
for wav_path in tqdm(recordings_path_set, desc="Submitting jobs"):
futures.append(ex.submit(_parse_recording, wav_path))
for future in tqdm(futures, desc="Processing recordings"):
try:
recording, recording_id = future.result()
recording_dict[recording_id] = recording
except Exception as e:
logging.warning(
f"Error processing recording {recording_id} with error: {e}"
)
recording_set = RecordingSet.from_recordings(recording_dict.values())
logging.info("Starting to process supervisions...")
# Step 3: Process supervisions
supervisions = []
for supervision in tqdm(supervision_list, desc="Processing supervisions"):
seg = _parse_supervision(supervision, recording_dict)
if seg is not None:
supervisions.append(seg)
logging.info("Processing Cuts...")
# Step 4: Create and validate manifests
supervision_set = SupervisionSet.from_segments(supervisions)
recording_set, supervision_set = fix_manifests(recording_set, supervision_set)
validate_recordings_and_supervisions(recording_set, supervision_set)
cut_set = CutSet.from_manifests(
recordings=recording_set, supervisions=supervision_set
)
cut_set = cut_set.sort_by_recording_id()
cut_set = cut_set.resample(sampling_rate)
cut_set = cut_set.trim_to_supervisions(keep_overlapping=False)
logging.info(f"Saving file to {output_dir / file_name}")
# Step 5: Write manifests to disk
cut_set.to_file(output_dir / file_name)
logging.info("Done!")
if __name__ == "__main__":
formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
logging.basicConfig(format=formatter, level=logging.INFO, force=True)
args = get_args()
prepare_dataset(
tsv_path=args.tsv_path,
prefix=args.prefix,
subset=args.subset,
sampling_rate=args.sampling_rate,
num_jobs=args.num_jobs,
output_dir=args.output_dir,
)
View Git Blame Copy Permalink