Improve usage

get_spectrogram.py

@@ -16,10 +16,9 @@ def arg_directory(path):
     else:
         raise argparse.ArgumentTypeError(f'`{path}` is not a valid path')
 
-def create_spectrogram(y, directory, filename, offset, duration):
-    window_size = 1024
-    window = np.hanning(window_size)
-    stft = librosa.core.spectrum.stft(y, n_fft=window_size, hop_length=512, window=window)
+def create_spectrogram(y, directory, filename, offset, duration, window_arg, hop_length_arg):
+    window = np.hanning(window_arg)
+    stft = librosa.core.spectrum.stft(y, n_fft=window_arg, hop_length=hop_length_arg, window=window)
 
     plt.close()
     plt.figure()
@@ -40,31 +39,40 @@ def create_spectrogram(y, directory, filename, offset, duration):
 
 def process_recordings(args):
     _, (i) = args
-    duration = 8
-    overlap = 2
+    duration = args.duration
+    overlap = args.overlap
 
     for count in range(args.img_per_rec):
         offset = count * (duration - overlap)
         filename = str(i[0])
         
         try:
-            y, _ = librosa.load(filename, offset=offset, duration=duration, sr=None)
-            create_spectrogram(y, args.directory, filename, offset, duration)
+            y, _ = librosa.load(filename, offset=offset, duration=duration, sr=args.sr)
+            create_spectrogram(y, args.directory, filename, offset, duration, args.window, args.hop)
         except Exception:
-            print(filename)
+            print(f'`{filename}` cannot be open...')
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, description='TODO')
-    parser.add_argument('-f', '--file', type=str, help='Name of the file that contains the recording to print')
-    parser.add_argument('-p', '--path_to_data', type=arg_directory, help='Path of the folder that contains the recordings', required=True)
-    parser.add_argument('-d', '--directory', type=arg_directory, help='Directory to which spectrograms will be stored', required=True)
-    parser.add_argument('-m', '--mode', type=str, choices=['unique', 'multiple'], help='Direction of the saved spectrogram')
-    parser.add_argument('-n', '--columns_name', type=str, help='Name of the columns that contain the path of the .wav')
-    parser.add_argument('-i', '--input', type=str, choices=['file', 'folder'], help='Choose "file" if you have a .csv file or "folder" to export spectrogram from all the .wav of a folder')
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, description='Extract spectrogram for each .wav file')
+    parser.add_argument('path_to_data', type=arg_directory, help='Path of the folder that contains the recordings', required=True)
+    parser.add_argument('directory', type=arg_directory, help='Directory to which spectrograms will be stored', required=True)
+    parser.add_argument('-m', '--mode', type=str, choices=['unique', 'multiple'], help='Direction of the saved spectrogram', default='multiple')
+    parser.add_argument('-n', '--columns_name', type=str, help='Name of the columns that contain the path of the .wav', default='Path')
+    parser.add_argument('-i', '--input', type=str, choices=['file', 'folder'], help='Choose "file" if you have a .csv file or "folder" to export '
+                                                                            'spectrogram from all the .wav of a folder', default='folder')
+    parser.add_argument('-f', '--file', type=str, help='Name of the file that contains the recording to print', default=None)
+    parser.add_argument('--frames', type=int, help='Number of spectrogram per file', default=30)
+    parser.add_argument('--duration', type=int, help='Duration for each spectrogram', default=8)
+    parser.add_argument('--overlap', type=int, help='Overlap between 2 spectrograms', default=2)
+    parser.add_argument('--sr', type=int, help='Sampling rate for the spectrogram. If no argument, '
+                                            'SR will be original SR of the recording', default=None)
+    parser.add_argument('--window', type=int, help='Window size for the Fourier Transform', default=1024)
+    parser.add_argument('--hop', type=int, help='Hop lenght for the Fourier Transform', default=512)
+    parser.add_argument('--cpu', type=int, help='To speed up the process, write 2 or more', default=1)
     args = parser.parse_args()
 
     if args.mode == 'multiple':
-        img_per_rec = 30
+        img_per_rec = args.frames
     elif args.mode == 'unique':
         img_per_rec = 1
 
@@ -76,5 +84,5 @@ if __name__ == "__main__":
     elif args.input == 'folder':
         df = pd.DataFrame(glob.glob(os.path.join(path_to_data, '*'), recursive=True), columns=['Path'])
 
-    p_map(process_recordings, enumerate(df.groupby('Path'), img_per_rec=img_per_rec), num_cpus=1, total=len(df.groupby('Path')))
+    p_map(process_recordings, enumerate(df.groupby('Path'), img_per_rec=img_per_rec), num_cpus=args.cpu, total=len(df.groupby('Path')))
     print(f'Saved to {args.directory}/Spectrogram')