JSIR-annotator/export.py at main · MALerLab/JSIR-annotator · 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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
from pathlib import Path
import shutil
import json
import argparse
from tqdm.auto import tqdm

DEFAULT_EXPORT_DIR = Path("/home/issyun/storage/JSIR")

parser = argparse.ArgumentParser()
parser.add_argument("export_dir", nargs="?", type=Path, default=DEFAULT_EXPORT_DIR)
parser.add_argument("--exclude-audio", action="store_true")
args = parser.parse_args()

EXPORT_DIR = args.export_dir

with open("metadata.json") as f:
  metadata = json.load(f)

metadata = [x for x in metadata if x["completed"] == True]

try:
  EXPORT_DIR.mkdir(parents=True, exist_ok=True)
  print(f"Exporting {len(metadata)} items to {EXPORT_DIR}")

  for item in metadata:
    if "sections" in item["files"]:
      del item["files"]["sections"]

  if not args.exclude_audio:
    print("Copying audio...")
    audio_dir = EXPORT_DIR / "audio"
    audio_dir.mkdir(parents=True, exist_ok=True)
    for item in tqdm(metadata, desc="Exporting"):
      src = Path(item["files"]["audio"])
      dst = audio_dir / src.name
      shutil.copy(src, dst)

  print("Copying labels...")
  beats_dir = EXPORT_DIR / "beats"
  chords_dir = EXPORT_DIR / "chords"
  structure_dir = EXPORT_DIR / "structure"
  segments_dir = EXPORT_DIR / "segments"
  for dir in [beats_dir, chords_dir, structure_dir, segments_dir]:
    dir.mkdir(parents=True, exist_ok=True)
    for file in dir.iterdir():
      if file.is_file():
        file.unlink()

  for item in tqdm(metadata, desc="Exporting"):
    if "beats" not in item["files"]:
      item["files"]["beats"] = f"beats/{Path(item['yt_id']).with_suffix('.txt')}"
    src = Path(item["files"]["beats"])
    dst = beats_dir / src.name
    shutil.copy(src, dst)

    if "chords" not in item["files"]:
      item["files"]["chords"] = f"chords/{Path(item['yt_id']).with_suffix('.csv')}"
    src = Path(item["files"]["chords"])
    dst = chords_dir / src.name
    shutil.copy(src, dst)

    if "structure" not in item["files"]:
      item["files"]["structure"] = f"structure/{Path(item['yt_id']).with_suffix('.csv')}"
    src = Path(item["files"]["structure"])
    dst = structure_dir / src.name
    shutil.copy(src, dst)

    if "segments" not in item["files"]:
      item["files"]["segments"] = f"segments/{Path(item['yt_id']).with_suffix('.csv')}"
    src = Path(item["files"]["segments"])
    dst = segments_dir / src.name
    shutil.copy(src, dst)

  if not args.exclude_audio:
    with open(EXPORT_DIR / "audio" / ".gitignore", "w") as f:
      f.write("*")

  print("Exporting metadata...")
  standards = set([x["standard"].lower() for x in metadata])
  with open("lead_sheets.json") as f:
    lead_sheets = json.load(f)
  lead_sheets = [x for x in lead_sheets if x["title"].lower() in standards]

  metadata.sort(key=lambda x: x["standard"].lower())
  lead_sheets.sort(key=lambda x: x["title"].lower())

  with open(EXPORT_DIR / "metadata.json", "w") as f:
    json.dump(metadata, f, indent=2)

  with open(EXPORT_DIR / "lead_sheets.json", "w") as f:
    json.dump(lead_sheets, f, indent=2)

except Exception as e:
  print(f"Error: {e}")
  print("Cleaning up...")
  beats_dir = EXPORT_DIR / "beats"
  chords_dir = EXPORT_DIR / "chords"
  structure_dir = EXPORT_DIR / "structure"
  segments_dir = EXPORT_DIR / "segments"
  for dir in [beats_dir, chords_dir, structure_dir, segments_dir]:
    if dir.exists():
      for file in dir.iterdir():
        if file.is_file():
          file.unlink()