Updated database_stats.py

i just added more features. i talk about the updates in the readme anyways

Author: UtaUtaUtau

README.md

@@ -25,21 +25,25 @@ Update 06/15/2022: Added a feature that counts how many instances of each phonem
 This script is for getting some statistics for your NNSVS/ENUNU database. It reads all the USTs and LABs of the database (it will only count the LABs that already have USTs by default). You can drag and drop the database folder over the script like with lab2audacity, but this script also takes in additional arguments for settings. Here's the whole help documentation of the script, but this may be invoked with `database_stats.py -h` as well.
 
 ```
-usage: database_stats.py [-h] [--all-labs] [--include-pau] db
+usage: database_stats.py [-h] [--all-labs] [--include-pau] [--skip-diphone] [--write-diphone] db
 
 Calculates some statistics for your existing NNSVS database
 
 positional arguments:
-  db                 The database's folder location
+  db                   The database's folder location
 
 optional arguments:
-  -h, --help         show this help message and exit
-  --all-labs, -L     Include all LABs in the LAB pass.
-  --include-pau, -P  Include pau phoneme in the phoneme tally passes.
+  -h, --help           show this help message and exit
+  --all-labs, -l       Include all LABs in the LAB pass.
+  --include-pau, -p    Include pau phoneme in the phoneme tally passes.
+  --skip-diphone, -s   Skip diphone density in calculations.
+  --write-diphone, -w  Include diphone density in the .csv version.
 ```
 
 This script will then save a `stats.txt` file inside the database folder you have passed to it.
 
+**Update 09/27/2022:** This script will now generate a `stats.csv` file as well to import statistics into a spreadsheet format automatically.
+
 ### lab2ust
 The whole lab2ust folder will be put in the UTAU plugins folder.
 

database_stats/database_stats.py

@@ -4,6 +4,7 @@
 import sys
 import time
 import traceback
+import csv
 from argparse import ArgumentParser
 
 notes = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
@@ -19,11 +20,14 @@ def midi_to_hz(x):
 try:
     parser = ArgumentParser(description='Calculates some statistics for your existing NNSVS database')
     parser.add_argument('db', help="The database's folder location")
-    parser.add_argument('--all-labs', '-L', action='store_true', help="Include all LABs in the LAB pass.")
-    parser.add_argument('--include-pau', '-P', action='store_true', help="Include pau phoneme in the phoneme tally passes.")
+    parser.add_argument('--all-labs', '-l', action='store_true', help="Include all LABs in the LAB pass.")
+    parser.add_argument('--include-pau', '-p', action='store_true', help="Include pau phoneme in the phoneme tally passes.")
+    parser.add_argument('--skip-diphone', '-s', action='store_true', help="Skip diphone density in calculations.")
+    parser.add_argument('--write-diphone', '-w', action='store_true', help="Include diphone density in the .csv version.")
 
     args, _ = parser.parse_known_args()
     db = args.db
+    calc_diphone = not args.skip_diphone
 
     print('Finding all USTs . . . ')
     usts = glob.glob(db + '/**/*.ust', recursive=True)
@@ -38,7 +42,7 @@ def midi_to_hz(x):
     else:
         labs = glob.glob(db + '/**/*.lab', recursive=True)
 
-    print('Caching LABs . . . ')
+    print('Reading LABs . . . ')
     phones = []
     for i in labs:
         temp = []
@@ -59,22 +63,24 @@ def midi_to_hz(x):
 
     mono_dens = dict(sorted(mono_dens.items(), key=lambda x : x[1], reverse=True))
 
-    print('Tallying diphones . . . ')
-    diph_dens = {}
-    for i in phones:
-        N = len(i)
-        for j in range(0, N-2):
-            k = i[j] + ' ' + i[j+1]
-            if (i[j] != 'pau' and i[j+1] != 'pau') or args.include_pau:
-                try:
-                    diph_dens[k] += 1
-                except KeyError:
-                    diph_dens[k] = 1
-
-    diph_dens = dict(sorted(diph_dens.items(), key=lambda x : x[1], reverse=True))
+    if calc_diphone:
+        print('Tallying diphones . . . ')
+        diph_dens = {}
+        for i in phones:
+            N = len(i)
+            for j in range(0, N-2):
+                k = i[j] + ' ' + i[j+1]
+                if (i[j] != 'pau' and i[j+1] != 'pau') or args.include_pau:
+                    try:
+                        diph_dens[k] += 1
+                    except KeyError:
+                        diph_dens[k] = 1
+
+        diph_dens = dict(sorted(diph_dens.items(), key=lambda x : x[1], reverse=True))
 
     print('Tallying notes . . . ')
-    note_dens = {}
+    note_dens_length = {}
+    note_dens_presence = {}
     note_count = 0
     for i in usts:
         ust = pyutau.UtauPlugin(i)
@@ -83,44 +89,104 @@ def midi_to_hz(x):
             if note.lyric not in ['R', 'pau']:
                 note_count += 1
                 try:
-                    note_dens[note.note_num] += note.length
+                    note_dens_length[note.note_num] += note.length
                 except KeyError:
-                    note_dens[note.note_num] = note.length
+                    note_dens_length[note.note_num] = note.length
+
+                try:
+                    note_dens_presence[note.note_num] += 1
+                except KeyError:
+                    note_dens_presence[note.note_num] = 1
 
     print('Calculating things idk . . . ')
-    note_dens = dict(sorted(note_dens.items(), reverse=True))
-    total_note = sum(note_dens.values())
+    note_dens_length = dict(sorted(note_dens_length.items(), reverse=True))
+    note_dens_presence = dict(sorted(note_dens_presence.items(), reverse=True))
+    total_note_length = sum(note_dens_length.values())
     mean_note = 0
-    for k, v in note_dens.items():
+    for k, v in note_dens_length.items():
         mean_note += k * v
-    mean_note /= total_note
+    mean_note /= total_note_length
 
-    print('Writing out stats in a text file . . . ')
-    note_range = list(note_dens.keys())
+    note_range = list(note_dens_length.keys())
     lo_range = midi_to_note(note_range[-1])
     hi_range = midi_to_note(note_range[0])
     lo_hz = midi_to_hz(note_range[-1])
     hi_hz = midi_to_hz(note_range[0])
     closest_note = midi_to_note(int(round(mean_note)))
     mean_hz = midi_to_hz(mean_note)
-    
+
+    basic_info = {
+            'Overall Range' : f'{lo_range} ~ {hi_range} ({lo_hz:.3f} ~ {hi_hz:.3f})',
+            'Note Count' : note_count,
+            'Total Note Length (UTAU length)' : total_note_length,
+            'Average Pitch' : f'{mean_hz:.3f} Hz (~{closest_note})'
+        }
+
+    print('Writing out stats in a text file . . . ')
     with open(db + '/stats.txt', 'w') as f:
-        f.write(f'Overall Range: {lo_range} ~ {hi_range} ({lo_hz:.3f} ~ {hi_hz:.3f})\n')
-        f.write(f'Note Count: {note_count}\n')
-        f.write(f'Total Note Length (UTAU length): {total_note}\n')
-        f.write(f'Average Pitch: {mean_hz:.3f} Hz (~{closest_note})\n\n')
-        
-        f.write(f'Note Density (based on UTAU length now YAY)\n')
-        for k, v in note_dens.items():
+        for k, v in basic_info.items():
+            f.write(f'{k}: {str(v)}\n')
+
+        f.write(f'\nNote Density (based on note length)\n')
+        for k, v in note_dens_length.items():
+            f.write(f'{midi_to_note(k)}: {v}\n')
+
+        f.write(f'\nNote Density (based on note presence)\n')
+        for k, v in note_dens_presence.items():
             f.write(f'{midi_to_note(k)}: {v}\n')
 
         f.write(f'\nMonophone Density\n')
         for k, v in mono_dens.items():
             f.write(f'{k}: {v}\n')
 
-        f.write(f'\nDiphone Density\n')
-        for k, v in diph_dens.items():
-            f.write(f'{k}: {v}\n')
+        if calc_diphone:
+            f.write(f'\nDiphone Density\n')
+            for k, v in diph_dens.items():
+                f.write(f'{k}: {v}\n')
+
+    print('Writing out stats in a .csv file . . . ')
+    header = ['Note', 'Density (lengths)', 'Density (presence)', '', 'Phoneme', 'Density', '', '', '']
+
+    if calc_diphone and args.write_diphone:
+        header = header[:7] + ['Diphone', 'Density', ''] + header[7:]
+
+    cols = len(header)
+    rows = max(len(note_dens_length), len(mono_dens)) + 1
+
+    if calc_diphone and args.write_diphone:
+        rows = max(len(note_dens_length), len(mono_dens), len(diph_dens)) + 1
+
+    sheet = [['' for c in range(cols)] for r in range(rows)]
+
+    sheet[0] = header
+    for r in range(1, rows):
+        i = r - 1
+        if i < len(note_dens_length):
+            dens_len = list(note_dens_length.items())[i]
+            dens_pres = list(note_dens_presence.values())[i]
+            sheet[r][0] = midi_to_note(dens_len[0])
+            sheet[r][1] = dens_len[1]
+            sheet[r][2] = dens_pres
+
+        if i < len(mono_dens):
+            dens = list(mono_dens.items())[i]
+            sheet[r][4] = dens[0]
+            sheet[r][5] = dens[1]
+
+        if calc_diphone and args.write_diphone:
+            if i < len(diph_dens):
+                dens = list(diph_dens.items())[i]
+                sheet[r][7] = dens[0]
+                sheet[r][8] = dens[1]
+
+        if i < len(basic_info):
+            info = list(basic_info.items())[i]
+            sheet[r][-2] = info[0]
+            sheet[r][-1] = info[1]
+
+    with open(db + '/stats.csv', 'w', newline='') as f:
+        writer = csv.writer(f)
+        writer.writerows(sheet)
 
 except Exception as e:
     for i in traceback.format_exception(e.__class__, e, e.__traceback__):