Initial commit
This commit is contained in:
Yatharth Sharma
274
zipvoice/bin/prepare_dataset.py
Normal file
274
zipvoice/bin/prepare_dataset.py
Normal file
@@ -0,0 +1,274 @@
|
||||
#!/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,
|
||||
)
|
||||
Reference in New Issue
Block a user