xjsc0's picture
1
64ec292
import argparse
import os
import time
from typing import Dict, Union
import soundfile as sf
def parse_args(dict_args: Union[Dict, None]) -> argparse.Namespace:
"""
Parse command-line arguments for configuring the model, dataset, and training parameters.
Args:
dict_args: Dict of command-line arguments. If None, arguments will be parsed from sys.argv.
Returns:
Namespace object containing parsed arguments and their values.
"""
parser = argparse.ArgumentParser()
parser.add_argument(
"--input_directory", type=str, help="Path to the input directory."
)
parser.add_argument(
"--output_directory", type=str, help="Path to the output directory."
)
parser.add_argument("--start_sec", type=float, default=20.0)
parser.add_argument("--end_sec", type=float, default=30.0)
parser.add_argument("--codec", type=str, default="wav")
parser.add_argument(
"--max_folders",
type=int,
default=float("inf"),
help="Maximum number of folders to process.",
)
if dict_args is not None:
args = parser.parse_args([])
args_dict = vars(args)
args_dict.update(dict_args)
args = argparse.Namespace(**args_dict)
else:
args = parser.parse_args()
if not args.output_directory:
original_dir = os.path.dirname(args.input_directory)
tests_dir = os.path.join("tests", original_dir)
os.makedirs(tests_dir, exist_ok=True)
args.output_directory = os.path.join(
tests_dir, os.path.basename(args.input_directory)
)
return args
def trim_wav(
input_file: str, output_file: str, start_sec: float, end_sec: float, codec: str
):
data, samplerate = sf.read(input_file)
start_sample = int(start_sec * samplerate)
end_sample = int(end_sec * samplerate)
trimmed_data = data[start_sample:end_sample]
sf.write(output_file, trimmed_data, samplerate, format=codec)
def trim_directory(args):
args = parse_args(args)
input_directory = args.input_directory
output_directory = args.output_directory
start_sec = args.start_sec
end_sec = args.end_sec
codec = args.codec
max_folder = args.max_folders
folder_count = 0
start_time = time.time()
for root, dirs, files in os.walk(input_directory):
if folder_count >= max_folder:
break
if os.path.relpath(root, input_directory) != ".":
print(f"Processing folder: {os.path.relpath(root, input_directory)}")
relative_path = os.path.relpath(root, input_directory)
target_folder = os.path.join(output_directory, relative_path)
os.makedirs(target_folder, exist_ok=True)
for filename in files:
if filename.endswith(f".{codec}"):
input_file = os.path.join(root, filename)
output_file = os.path.join(
target_folder, filename.replace(f".{codec}", f".{codec}")
)
try:
trim_wav(input_file, output_file, start_sec, end_sec, codec)
except Exception as e:
print(f"Error processing {filename} in folder {root}: {e}")
folder_count += 1
end_time = time.time()
total_time = end_time - start_time
print(f"Processing complete. Total time: {total_time:.2f} seconds.")
return output_directory
if __name__ == "__main__":
trim_directory(None)