bcresnet/split.py at main · rolyantrauts/bcresnet · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import argparse
import glob
import os
import shutil

def split(source_dir='./in', dest_dir='./out', testing_percent=0.05, validation_percent=0.15, label="kw"):

  if not os.path.exists(dest_dir):
    os.makedirs(dest_dir)
  if not os.path.exists(dest_dir + "/training"):
    os.makedirs(dest_dir + "/training")
  if not os.path.exists(dest_dir + "/testing"):
    os.makedirs(dest_dir + "/testing")
  if not os.path.exists(dest_dir + "/validation"):
    os.makedirs(dest_dir + "/validation")
  if not os.path.exists(dest_dir + "/training/" + label):
    os.makedirs(dest_dir + "/training/" + label)
  if not os.path.exists(dest_dir + "/testing/" + label):
    os.makedirs(dest_dir + "/testing/" + label)
  if not os.path.exists(dest_dir + "/validation/" + label):
    os.makedirs(dest_dir + "/validation/" + label)

  if not os.path.exists(source_dir):
    print("Source_dir = " + source_dir + " does not exist!")
    exit()

  source = glob.glob(os.path.join(source_dir, '*.wav'))
  if len(source) < 2:
    print("No files found *.wav")
    exit()
  count = 0
  testing_qty = len(source) * testing_percent
  validation_qty = len(source) * validation_percent
  for wav in source:
    if count < testing_qty:
      shutil.move(wav, dest_dir + "/testing/" + label + "/" + os.path.basename(wav))
    elif count < validation_qty + testing_qty:
      shutil.move(wav, dest_dir + "/validation/" + label + "/" + os.path.basename(wav))
    else:
      shutil.move(wav, dest_dir + "/training/" + label + "/" + os.path.basename(wav))
    count += 1


def main_body():
  parser = argparse.ArgumentParser()
  parser.add_argument('--source_dir', default='./in', help='source dir location')
  parser.add_argument('--dest_dir', type=str, default='./out', help='dest dir location')
  parser.add_argument('--testing_percent', type=float, default=0.05, help='decimal percentage of testing qty')
  parser.add_argument('--validation_percent', type=float, default=0.15, help='decimal percentage of validation qty')
  parser.add_argument('--label', type=str, default='kw', help='label name')
  args = parser.parse_args()

  if args.dest_dir == None:
    args.dest_dir = "./out"

  split(args.source_dir, args.dest_dir, args.testing_percent, args.validation_percent, args.label)

if __name__ == '__main__':
  main_body()