Skip to content

Commit

Permalink
Add output_dir argument to csv_to_wfdb. Fixes #67.
Browse files Browse the repository at this point in the history
  • Loading branch information
tompollard committed Jul 1, 2024
1 parent 34b989e commit d6d5aa7
Showing 1 changed file with 22 additions and 8 deletions.
30 changes: 22 additions & 8 deletions wfdb/io/convert/csv.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ def csv_to_wfdb(
header=True,
delimiter=",",
verbose=False,
output_dir=None
):
"""
Read a WFDB header file and return either a `Record` object with the
Expand Down Expand Up @@ -235,6 +236,9 @@ def csv_to_wfdb(
verbose : bool, optional
Whether to print all the information read about the file (True) or
not (False).
output_dir : str, optional
The directory where the output files will be saved. If not provided,
the output files will be saved in the same directory as the input file.
Returns
-------
Expand Down Expand Up @@ -291,6 +295,7 @@ def csv_to_wfdb(
df_CSV = pd.read_csv(file_name, delimiter=delimiter, header=None)
if verbose:
print("Successfully read CSV")

# Extract the entire signal from the dataframe
p_signal = df_CSV.values
# The dataframe should be in (`sig_len`, `n_sig`) dimensions
Expand All @@ -300,6 +305,7 @@ def csv_to_wfdb(
n_sig = p_signal.shape[1]
if verbose:
print("Number of signals: {}".format(n_sig))

# Check if signal names are valid and set defaults
if not sig_name:
if header:
Expand All @@ -318,15 +324,23 @@ def csv_to_wfdb(
if verbose:
print("Signal names: {}".format(sig_name))

# Set the output header file name to be the same, remove path
if os.sep in file_name:
file_name = file_name.split(os.sep)[-1]
record_name = file_name.replace(".csv", "")
# Determine the output directory
if output_dir:
if not os.path.exists(output_dir):
os.makedirs(output_dir)
output_base = os.path.join(
output_dir, os.path.basename(file_name).replace(".csv", "")
)
else:
if os.sep in file_name:
file_name = file_name.split(os.sep)[-1]
output_base = file_name.replace(".csv", "")

if verbose:
print("Output header: {}.hea".format(record_name))
print("Output base: {}".format(output_base))

# Replace the CSV file tag with DAT
dat_file_name = file_name.replace(".csv", ".dat")
dat_file_name = output_base + ".dat"
dat_file_name = [dat_file_name] * n_sig
if verbose:
print("Output record: {}".format(dat_file_name[0]))
Expand Down Expand Up @@ -419,7 +433,7 @@ def csv_to_wfdb(
if record_only:
# Create the record from the input and generated values
record = Record(
record_name=record_name,
record_name=output_base,
n_sig=n_sig,
fs=fs,
samps_per_frame=samps_per_frame,
Expand Down Expand Up @@ -454,7 +468,7 @@ def csv_to_wfdb(
else:
# Write the information to a record and header file
wrsamp(
record_name=record_name,
record_name=output_base,
fs=fs,
units=units,
sig_name=sig_name,
Expand Down

0 comments on commit d6d5aa7

Please sign in to comment.