Spaces:
Running
on
Zero
Running
on
Zero
File size: 8,549 Bytes
82bc972 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 |
# for each each audio segment, find the non-overlapping neighboring segments
from importlib.resources import path
import pathlib
import soundfile as sf
import numpy as np
import json
import multiprocessing
import argparse
import tqdm
import gzip
import time
import os
from tokenizer import TextTokenizer, tokenize_text
import glob
import sys
import os, random, numpy as np, socket
import json
import tqdm
import json
import tqdm
def write_jsonl(data, fn):
with open(fn, "w") as file:
for entry in data:
file.write(json.dumps(entry, ensure_ascii=False) + "\n")
def read_jsonl(file_path):
cur_data = []
with open(file_path, 'r', encoding='utf-8-sig') as file:
for line in file:
cur_data.append(json.loads(line.strip()))
return cur_data
from collections import defaultdict
# Function to create a defaultdict recursively
def nested_defaultdict(levels, inner_type):
if levels <= 1:
return defaultdict(inner_type)
return defaultdict(lambda: nested_defaultdict(levels-1, inner_type))
def find_neighbor(args):
split2manifest = {
"train": [
"libriheavy_cuts_small.jsonl",
"libriheavy_cuts_medium.jsonl",
"libriheavy_cuts_large.jsonl",
"libriheavy_long_cuts_small.jsonl",
"libriheavy_long_cuts_medium.jsonl",
"libriheavy_long_cuts_large.jsonl"
],
"valid": [
"libriheavy_cuts_dev.jsonl",
"libriheavy_long_cuts_dev.jsonl"
],
"test": [
"libriheavy_cuts_test_clean.jsonl",
"libriheavy_cuts_test_other.jsonl",
"libriheavy_long_cuts_test_clean.jsonl",
"libriheavy_long_cuts_test_other.jsonl"
]
}
stime = time.time()
organized_data = nested_defaultdict(4, list)
for mani_fn in split2manifest[args.split]:
# data = open_mani(os.path.join(mani_dir, mani_fn))
mani_full_fn = os.path.join(args.manifest_dir, mani_fn)
data = read_jsonl(mani_full_fn)
for item in data:
file_id = item['supervisions'][0]['id'] + '.flac'
recording_id = item['recording']['id'] + '.flac'
sizeSplit, spk, book, flac = recording_id.split("/") # e.g. 'medium/100/emerald_city_librivox_64kb_mp3/emeraldcity_01_baum_64kb'
if os.path.isfile(os.path.join(args.audio_dir, recording_id)):
vad = (item['start'], item['start']+item['duration'])
text = item['supervisions'][0]['custom']['texts'][0]
file_id = file_id.replace(".flac", "") + f"_{vad[0]:.2f}_{vad[1]:.2f}.flac"
organized_data[sizeSplit][spk][book][recording_id].append({"file_id": file_id, "vad":vad, "text": text})
# # for each recording_id, find the non-overlapping neighboring segments based on vad
# for sizeSplit in organized_data:
# for spk in organized_data[sizeSplit]:
# for book in organized_data[sizeSplit][spk]:
# for recording_id in organized_data[sizeSplit][spk][book]:
# segments = organized_data[sizeSplit][spk][book][recording_id]
# segments.sort(key=lambda x: x['vad'][0])
# for i in range(len(segments)):
# # for segment i, find the non-overlapping neighboring segments
# write_fn = os.path.join(args.output_dir, f"{segments[i]['file_id'].replace('.flac', '.txt')}")
# neighbors = []
# distance = []
# for j in range(len(segments)):
# if segments[i]['vad'][1] < segments[j]['vad'][0] or segments[i]['vad'][0] > segments[j]['vad'][0]:
# neighbors.append(segments[j]['file_id'].replace('.flac', '.txt'))
# distance.append(min(abs(segments[i]['vad'][1] - segments[j]['vad'][0]), abs(segments[i]['vad'][0] - segments[j]['vad'][1])))
# # order neighbors by distance
# neighbors_distance = [[x, dist] for dist, x in sorted(zip(distance, neighbors))]
# os.makedirs(os.path.dirname(write_fn), exist_ok=True)
# with open(write_fn, "w") as f:
# # note that there might be no neighbors, in which case the file is empty
# for neighbor, dist in neighbors_distance:
# f.write(f"{neighbor}\t{dist}\n")
# use multiprocessing.Pool for the above
segments = [organized_data[sizeSplit][spk][book][recording_id] for sizeSplit in organized_data for spk in organized_data[sizeSplit] for book in organized_data[sizeSplit][spk] for recording_id in organized_data[sizeSplit][spk][book]]
# only keep those that are exist
print(f"originally total {len(segments)} segments")
segments = [seg for seg in segments if os.path.isfile(os.path.join("/".join(args.output_dir.split("/")[:-1]),"audio", seg[0]['file_id']))]
print(f"after check existance, total {len(segments)} segments")
print(f"organizing took {(time.time()-stime)/60:.2f} minutes")
with multiprocessing.Pool(processes=args.n_workers) as pool:
for _ in tqdm.tqdm(pool.imap_unordered(find_neighbor_each, segments), total=len(segments)):
pass
# audio_root = "/data/scratch/pyp/datasets/librilight/preprocessed/audio"
def find_neighbor_each(segments):
# for each recording_id, find the non-overlapping neighboring segments based on vad
# only keep segments that have audio files
# actually only keep segments that have ipa_alignment files
segments = [seg for seg in segments if os.path.isfile(os.path.join("/".join(args.output_dir.split("/")[:-1]),"ipa_alignment", seg['file_id'].replace(".flac", ".txt")))]
if len(segments) <= 1:
return
for i in range(len(segments)):
# for segment i, find the non-overlapping neighboring segments
write_fn = os.path.join(args.output_dir, f"{segments[i]['file_id'].replace('.flac', '.txt')}")
neighbors = []
distance = []
for j in range(len(segments)):
if segments[i]['vad'][1] < segments[j]['vad'][0] or segments[i]['vad'][0] > segments[j]['vad'][0]:
neighbors.append(segments[j])
distance.append(min(abs(segments[i]['vad'][1] - segments[j]['vad'][0]), abs(segments[i]['vad'][0] - segments[j]['vad'][1])))
if len(neighbors) == 0:
continue
# order neighbors by distance
index = np.argsort(distance)
neighbors_distance = [[neighbors[ind], distance[ind]] for ind in index]
os.makedirs(os.path.dirname(write_fn), exist_ok=True)
with open(write_fn, "w") as f:
# note that there might be no neighbors, in which case the file is empty
for neighbor, dist in neighbors_distance:
f.write(f"{neighbor['file_id'].replace('.flac', '.txt')}\t{dist}\t{neighbor['vad'][1] - neighbor['vad'][0]}\n") # file_id, distance, duration
def parse_args():
parser = argparse.ArgumentParser(description="Cut a dataset in small "
"sequences using VAD files")
parser.add_argument('--split', type=str, default='train', choices=['train', 'valid', 'test'], help="train = libriheavy_cuts_{small,medium,large}.jsonl.gz, valid = libriheavy_cuts_dev_{clean,other}.jsonl.gz, test = libriheavy_cuts_test_{clean,other}.jsonl.gz")
parser.add_argument('--audio_dir', type=str, default="/data/scratch/pyp/datasets/librilight_example",
help="Path to the audio directory")
parser.add_argument('--manifest_dir', type=str, default="/data/scratch/pyp/datasets/librilight/libriheavy", help="path to the transcription file's dir, can be downloaded https://huggingface.co/datasets/pkufool/libriheavy/tree/main/v0.1")
parser.add_argument('--output_dir', type=str, default="/data/scratch/pyp/datasets/librilight/librilight_example_preprocessed/neighbors",
help="Path to the output directory")
parser.add_argument('--n_workers', type=int, default=16,
help="Number of parallel worker processes")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
pathlib.Path(args.output_dir).mkdir(exist_ok=True, parents=True)
find_neighbor(args) |