diff --git a/data/generate_dataset.py b/data/generate_dataset.py
index 242c652..6f5cb0c 100644
--- a/data/generate_dataset.py
+++ b/data/generate_dataset.py
@@ -42,9 +42,9 @@ def load_data(folder):
 
     fish_freq = np.load(base_path / 'analysis' / 'fish_freq.npy')
     rise_idx = np.load(base_path / 'analysis' / 'rise_idx.npy')
+    rise_size = np.load(base_path / 'analysis' / 'rise_size.npy')
 
-
-    return fill_freqs, fill_times, fill_spec, EODf_v, ident_v, idx_v, times_v, fish_freq, rise_idx
+    return fill_freqs, fill_times, fill_spec, EODf_v, ident_v, idx_v, times_v, fish_freq, rise_idx, rise_size
 
 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')
@@ -68,7 +68,7 @@ def main(args):
     d_time = 60*15
     time_overlap = 60*5
 
-    freq, times, spec, EODf_v, ident_v, idx_v, times_v, fish_freq, rise_idx = load_data(args.folder)
+    freq, times, spec, EODf_v, ident_v, idx_v, times_v, fish_freq, rise_idx, rise_size = 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)])
@@ -117,7 +117,10 @@ def main(args):
             times_v_idx0, times_v_idx1 = np.argmin(np.abs(times_v - t0)), np.argmin(np.abs(times_v - t1))
             for id_idx in range(len(fish_freq)):
                 ax.plot(times_v[times_v_idx0:times_v_idx1], fish_freq[id_idx][times_v_idx0:times_v_idx1], marker='.', color='k', markersize=4)
-                rise_idx_oi = np.array(rise_idx[id_idx][(rise_idx[id_idx] >= times_v_idx0) & (rise_idx[id_idx] <= times_v_idx1)], dtype=int)
+                rise_idx_oi = np.array(rise_idx[id_idx][
+                                           (rise_idx[id_idx] >= times_v_idx0) &
+                                           (rise_idx[id_idx] <= times_v_idx1) &
+                                           (rise_size[id_idx] >= 10)], dtype=int)
                 ax.plot(times_v[rise_idx_oi], fish_freq[id_idx][rise_idx_oi], 'o', color='tab:red')
 
             plt.show()