Change code layout to be usable by outside scripts

ipi_extract.py

@@ -396,16 +396,20 @@ class Callback(object):
         self.fax.get_figure().set_constrained_layout(False)
 
 
-def main(args):
-    print(f'Loading and processing {args.input}')
-    song, sr = read(args.input, always_2d=True)
-    song = song[:, args.channel]
-    sos = sg.butter(3, [2e3,20e3], 'bandpass', fs=sr, output='sos')
+def load_file(in_path, channel, low, high):
+    print(f'Loading and processing {in_path}')
+    song, sr = read(in_path, always_2d=True)
+    song = song[:, channel]
+    sos = sg.butter(3, [low, high], 'bandpass', fs=sr, output='sos')
     song = sg.sosfiltfilt(sos, song)
     frac = Fraction(FSSR, sr)
     song_resample = sg.resample_poly(song, frac.numerator, frac.denominator)
     print('Done processing')
+    return song, sr, song_resample
+
 
+def init(in_path, channel, low=2e3, high=20e3):
+    song, sr, song_resample = load_file(in_path, channel, low, high)
     fig = plt.figure('IPI', figsize=[16, 9], constrained_layout=True)
     gs = fig.add_gridspec(12, 20)
 
@@ -472,7 +476,8 @@ def main(args):
     callback.cursor = m_cursor
     for i in range(3):
         data_view[i][0][1] = (ax_view[i][0].axvline(10, c='k', linestyle='--'), ax_view[i][0].axvline(10, c='k'))
-        data_view[i][0][0] = ax_view[i][0].plot(np.linspace(0,20,int(20e-3*sr), False), np.zeros(int(20e-3*sr)))[0]
+        data_view[i][0][0] = ax_view[i][0].plot(np.linspace(0, 20, int(20e-3 * sr), False), np.zeros(int(20e-3 * sr)))[
+            0]
         ax_view[i][0].set_xlim(0, 20)
         ax_view[i][0].set_xlabel('IPI man:None')
         ax_view[i][0].set_ylim(-1, 1)
@@ -486,14 +491,16 @@ def main(args):
         data_view[i][1][1][1].set_visible(False)
 
         data_view[i][2][1] = ax_view[i][2].axvline(10, c='k')
-        data_view[i][2][0] = ax_view[i][2].plot(np.linspace(0, 10, int(10e-3*sr), False), np.zeros(int(10e-3*sr)))[0]
+        data_view[i][2][0] = ax_view[i][2].plot(np.linspace(0, 10, int(10e-3 * sr), False), np.zeros(int(10e-3 * sr)))[
+            0]
         ax_view[i][2].set_xlim(0, 10)
         ax_view[i][2].set_xlabel('IPI man:None auto:None')
         ax_view[i][2].set_ylim(-1, 1)
         data_view[i][2][1].set_visible(False)
 
         data_view[i][3][1] = ax_view[i][3].axvline(10, c='k')
-        data_view[i][3][0] = ax_view[i][3].plot(np.linspace(0, 10, int(10e-3*sr), False),np.zeros(int(10e-3*sr)))[0]
+        data_view[i][3][0] = ax_view[i][3].plot(np.linspace(0, 10, int(10e-3 * sr), False), np.zeros(int(10e-3 * sr)))[
+            0]
         ax_view[i][3].set_xlim(0, 10)
         ax_view[i][3].set_xlabel('IPI man:None auto:None')
         ax_view[i][3].set_ylim(0, 1)
@@ -514,24 +521,38 @@ def main(args):
     plt.draw()
     plt.pause(0.2)
     fig.set_constrained_layout(False)
+    return {'callback': callback, 'fig': fig, 'buttons':
+        {'b_left': b_left, 'b_right': b_right, 'play_b': play_b, 'resize_b': resize_b, 'r_button': r_button,
+         'fs_click': cid}}  # Needed to keep the callbacks alive
+
+
+def reset(in_path, channel, low=2e3, high=20e3):
+    song, sr, song_resample = load_file(in_path, channel, low, high)
+
+
+
+
+def main(args):
+    if args.out == '':
+        outpath = args.out.rsplit('.', 1)[0] + '.pred.h5'
+    else:
+        outpath = args.out
+    if os.path.isfile(outpath) and not args.erase:
+        return 1
+    ref_dict = init(args.path, args.channel)
     plt.show()
-    return pd.DataFrame.from_dict(callback.df, orient='index')
+    df = pd.DataFrame.from_dict(ref_dict['callback'].df, orient='index')
+    df.to_hdf(outpath, 'df')
+    return 0
 
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument("input", type=str, help="Input file")
-    parser.add_argument("--out", type=str, default='', help="Output file. Default to the input_path'.pred'")
+    parser.add_argument("--out", type=str, default='', help="Output file. Default to the input_path'.pred.h5'")
     parser.add_argument("--channel", type=int, default=0, help="Sound channel to be analysed. Indices start from 0.")
     parser.add_argument("--erase", action='store_true', help="If out file exist and this option is not given,"
                                                              " the computation will be halted")
 
     args = parser.parse_args()
-    if args.out == '':
-        outpath = args.out.rsplit('.', 1)[0] + '.pred'
-    else:
-        outpath = args.out
-    if os.path.isfile(outpath) and not args.erase:
-        quit()
-    df = main(args)
-    df.to_hdf(outpath, 'df')
+    main(args)