Adds output to sheetsage.infer

prepare.sh

@@ -3,6 +3,12 @@
 # NOTE: Override this with a local directory of your choosing
 SHEETSAGE_CACHE_DIR=$(python3 -c "import pathlib; print(pathlib.Path(pathlib.Path.home(), '.sheetsage').resolve())")
 
+if [ -f "$(pwd)/setup.py" ] && [ -d "$(pwd)/sheetsage" ]; then
+  DOCKER_LINK_LIB_ARG="-v $(pwd)/sheetsage:/sheetsage/sheetsage"
+else
+  DOCKER_LINK_LIB_ARG=""
+fi
+
 while [[ $# -gt 0 ]]; do
   case $1 in
     -j|--use_jukebox)
@@ -21,6 +27,7 @@ docker run \
   -it \
   --rm \
   -u $(id -u) \
+  ${DOCKER_LINK_LIB_ARG} \
   -v $SHEETSAGE_CACHE_DIR:/sheetsage/cache \
   chrisdonahue/sheetsage \
   /bin/bash -c \

sheetsage.sh

@@ -8,6 +8,12 @@ DOCKER_CPUS=$(python3 -c "import os; cpus=os.sched_getaffinity(0); print(','.joi
 DOCKER_CPU_ARG="--cpuset-cpus ${DOCKER_CPUS}"
 DOCKER_GPU_ARG=""
 
+if [ -f "$(pwd)/setup.py" ] && [ -d "$(pwd)/sheetsage" ]; then
+  DOCKER_LINK_LIB_ARG="-v $(pwd)/sheetsage:/sheetsage/sheetsage"
+else
+  DOCKER_LINK_LIB_ARG=""
+fi
+
 DOCKER_ARGS=()
 while [[ $# -gt 0 ]]; do
   case $1 in
@@ -57,6 +63,7 @@ docker run \
   ${DOCKER_CPU_ARG} \
   ${DOCKER_GPU_ARG} \
   -u $(id -u) \
+  ${DOCKER_LINK_LIB_ARG} \
   -v $SHEETSAGE_CACHE_DIR:/sheetsage/cache \
   -v $SHEETSAGE_OUTPUT_DIR:/sheetsage/output \
   chrisdonahue/sheetsage \

sheetsage/infer.py

@@ -712,6 +712,18 @@ def sheetsage(
         type=float,
         help="Approximate timestamp of end downbeat (to transcribe a segment of the audio).",
     )
+    parser.add_argument(
+        "-t",
+        "--title",
+        type=str,
+        help="Title of the song.",
+    )
+    parser.add_argument(
+        "-a",
+        "--artist",
+        type=str,
+        help="Name of the artist or composer.",
+    )
     parser.add_argument(
         "-o",
         "--output_dir",
@@ -767,6 +779,8 @@ def sheetsage(
     parser.set_defaults(
         segment_start_hint=None,
         segment_end_hint=None,
+        title=None,
+        artist=None,
         output_dir="./output",
         use_jukebox=False,
         measures_per_chunk=8,
@@ -781,7 +795,7 @@ def sheetsage(
 
     logging.basicConfig(level=logging.INFO)
 
-    lead_sheet = sheetsage(
+    lead_sheet, segment_beats, segment_beats_times = sheetsage(
         args.audio_path_or_url,
         segment_start_hint=args.segment_start_hint,
         segment_end_hint=args.segment_end_hint,
@@ -796,9 +810,31 @@ def sheetsage(
         tqdm=tqdm,
     )
 
+    # Create output directory
     output_dir = pathlib.Path(args.output_dir).resolve()
     if output_dir == pathlib.Path("./output").resolve():
         uuid = uuid.uuid4().hex
         output_dir = pathlib.Path(output_dir, uuid)
     logging.info(f"Writing to {output_dir}")
     output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Write lead sheet
+    lily = lead_sheet.as_lily(artist=args.artist, title=args.title)
+    with open(pathlib.Path(output_dir, "output.ly"), "w") as f:
+        f.write(lily)
+    with open(pathlib.Path(output_dir, "output.pdf"), "wb") as f:
+        f.write(
+            engrave(
+                lily, out_format="pdf", transparent=False, trim=False, hide_footer=False
+            )
+        )
+
+    # Write MIDI
+    with open(pathlib.Path(output_dir, "output.midi"), "wb") as f:
+        f.write(
+            lead_sheet.as_midi(
+                pulse_to_time_fn=create_beat_to_time_fn(
+                    segment_beats, segment_beats_times
+                )
+            )
+        )