diff --git a/wfdb/io/convert/csv.py b/wfdb/io/convert/csv.py index 3cfd25a2..42373401 100644 --- a/wfdb/io/convert/csv.py +++ b/wfdb/io/convert/csv.py @@ -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 @@ -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 ------- @@ -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 @@ -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: @@ -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])) @@ -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, @@ -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,