diff --git a/data/generate_dataset.py b/data/generate_dataset.py
index 1ee8e9f..7186688 100644
--- a/data/generate_dataset.py
+++ b/data/generate_dataset.py
@@ -1,4 +1,5 @@
 import numpy as np
+import argparse
 import torch
 from torch import nn
 import torch.nn.functional as F
@@ -40,14 +41,26 @@ def load_data(folder):
     times_v = np.load(base_path / 'times.npy')
 
     rise_idx = np.load(base_path / 'analysis' / 'rise_idx.npy')
-    ff = np.load(base_path / 'analysis' / 'fish_freq.npy')
-    embed()
-    quit()
 
     return fill_freqs, fill_times, fill_spec, EODf_v, ident_v, idx_v, times_v, rise_idx
 
+def save_spec_pic(folder, s_trans, times, freq, t_idx0, t_idx1, f_idx0, f_idx1, t_res, f_res):
+    fig_title = (f'{Path(folder).name}__{t0:.0f}s-{t1:.0f}s__{f0:4.0f}-{f1:4.0f}Hz').replace(' ', '0')
+    fig = plt.figure(figsize=(7, 7), num=fig_title)
+    gs = gridspec.GridSpec(1, 2, width_ratios=(8, 1), wspace=0)  # , bottom=0, left=0, right=1, top=1
+    gs2 = gridspec.GridSpec(1, 1, bottom=0, left=0, right=1, top=1)  #
+    ax = fig.add_subplot(gs2[0, 0])
+    im = ax.imshow(s_trans.squeeze(), cmap='gray', aspect='auto', origin='lower',
+                   extent=(times[t_idx0] / 3600, times[t_idx1] / 3600 + t_res, freq[f_idx0], freq[f_idx1] + f_res))
+    ax.axis(False)
+
+    plt.savefig(os.path.join('train', fig_title + '.png'), dpi=256)
+    plt.close()
+
+
+def main(args):
+
 
-def main(folder):
     min_freq = 200
     max_freq = 1500
     d_freq = 200
@@ -55,7 +68,7 @@ def main(folder):
     d_time = 60*15
     time_overlap = 60*5
 
-    freq, times, spec, EODf_v, ident_v, idx_v, times_v, rise_idx = load_data(folder)
+    freq, times, spec, EODf_v, ident_v, idx_v, times_v, rise_idx = load_data(args.folder)
     f_res, t_res = freq[1] - freq[0], times[1] - times[0]
 
     unique_ids = np.unique(ident_v[~np.isnan(ident_v)])
@@ -82,30 +95,33 @@ def main(folder):
         f_idx0, f_idx1 = np.argmin(np.abs(freq - f0)), np.argmin(np.abs(freq - f1))
         t_idx0, t_idx1 = np.argmin(np.abs(times - t0)), np.argmin(np.abs(times - t1))
 
-        embed()
-        quit()
-
         s = torch.from_numpy(spec[f_idx0:f_idx1, t_idx0:t_idx1].copy()).type(torch.float32)
         log_s = torch.log10(s)
         transformed = T.Normalize(mean=torch.mean(log_s), std=torch.std(log_s))
         s_trans = transformed(log_s.unsqueeze(0))
 
+        if not args.dev:
+            save_spec_pic(args.folder, s_trans, times, freq, t_idx0, t_idx1, f_idx0, f_idx1, t_res, f_res)
+            exit()
+        else:
+            fig_title = (f'{Path(folder).name}__{t0:.0f}s-{t1:.0f}s__{f0:4.0f}-{f1:4.0f}Hz').replace(' ', '0')
+            fig = plt.figure(figsize=(7, 7), num=fig_title)
+            gs = gridspec.GridSpec(1, 2, width_ratios=(8, 1), wspace=0)  # , bottom=0, left=0, right=1, top=1
+            ax = fig.add_subplot(gs[0, 0])
+            cax = fig.add_subplot(gs[0, 1])
+            im = ax.imshow(s_trans.squeeze(), cmap='gray', aspect='auto', origin='lower',
+                           extent=(times[t_idx0] / 3600, times[t_idx1] / 3600 + t_res, freq[f_idx0], freq[f_idx1] + f_res))
+
+            fig.colorbar(im, cax=cax, orientation='vertical')
+            plt.show()
 
-        fig_title = (f'{Path(folder).name}__{t0:.0f}s-{t1:.0f}s__{f0:4.0f}-{f1:4.0f}Hz').replace(' ', '0')
-        fig = plt.figure(figsize=(7, 7), num=fig_title)
-        gs = gridspec.GridSpec(1, 2, width_ratios=(8, 1), wspace=0)# , bottom=0, left=0, right=1, top=1
-        gs2 = gridspec.GridSpec(1, 1, bottom=0, left=0, right=1, top=1)#
-        ax = fig.add_subplot(gs2[0, 0])
-        im = ax.imshow(s_trans.squeeze(), cmap='gray', aspect='auto', origin='lower', extent=(times[t_idx0]/3600, times[t_idx1]/3600 + t_res, freq[f_idx0], freq[f_idx1] + f_res))
-        ax.axis(False)
-
-        plt.savefig(os.path.join('train', fig_title + '.png'), dpi=256)
-        plt.close()
     # # ax.imshow(spec[f0:f1, t0:t1], cmap='gray')
 
-
-
-
-
 if __name__ == '__main__':
-    main(sys.argv[1])
\ No newline at end of file
+    parser = argparse.ArgumentParser(description='Evaluated electrode array recordings with multiple fish.')
+    parser.add_argument('file', type=str, help='single recording analysis', default='')
+    parser.add_argument('-d', "--dev", action="store_true", help="developer mode; no data saved")
+    # parser.add_argument('-x', type=int, nargs=2, default=[1272, 1282], help='x-borders of LED detect area (in pixels)')
+    # parser.add_argument('-y', type=int, nargs=2, default=[1500, 1516], help='y-borders of LED area (in pixels)')
+    args = parser.parse_args()
+    main(args)
\ No newline at end of file