Add functioning viz tool

Author: sevagh

headbang/consensus.py

@@ -15,7 +15,7 @@
 
 
 algo_names = [
-    "_", # dummy at index 0 because in this code, beat trackers start at 1: "1,2,3...8"
+    "_",  # dummy at index 0 because in this code, beat trackers start at 1: "1,2,3...8"
     "madmom DBNBeatTrackingProcessor",
     "madmom BeatDetectionProcessor",
     "essentia BeatTrackerMultiFeature",

headbang/viz_tool.py

@@ -6,9 +6,11 @@
 import argparse
 import sys
 import librosa
+import pandas as pd
 import gc
 import os
 import multiprocessing
+from collections import OrderedDict
 from moviepy.editor import *
 from moviepy.audio.AudioClip import AudioArrayClip
 from tempfile import gettempdir
@@ -19,10 +21,18 @@
 from madmom.io.audio import write_wave_file
 
 
+def find_closest(A, target):
+    # A must be sorted
+    idx = A.searchsorted(target)
+    idx = numpy.clip(idx, 1, len(A) - 1)
+    left = A[idx - 1]
+    right = A[idx]
+    idx -= target - left < right - target
+    return idx
+
+
 def main():
-    parser = argparse.ArgumentParser(
-        description="Vizualize the headbang beat tracker"
-    )
+    parser = argparse.ArgumentParser(description="Vizualize the headbang beat tracker")
 
     parser.add_argument("wav_in", type=str, help="wav file to process")
     parser.add_argument("mp4_out", type=str, help="mp4 output path")
@@ -41,32 +51,17 @@ def main():
     # get the inner multi-beat-tracker list from headbangbeattracker's consensusbeattracker object
     individual_tracker_beat_locations = hbt.cbt.beat_results
 
-    # blue, yellow, magenta, violet, orange, brown, white
-    colors = itertools.cycle([(0, 165, 255), (255, 0, 255), (255, 255, 0), (255, 69, 0), (0, 255, 255), (165, 42, 42), (255, 255, 255)])
-
-    beat_trackers = {
-            'headbang': {
-                'beats': strong_beat_locations,
-                'color': (255, 0, 0), # red
-            },
-            'consensus': {
-                'beats': hbt.beat_consensus,
-                'color': (0, 255, 0), # lime green
-            },
-    }
-    for i, algo_name in enumerate(algo_names[1:]):
-        beat_trackers[algo_name] = {
-            'beats': individual_tracker_beat_locations[i],
-            'color': next(colors)
-        }
-
-    beat_trackers['onsets'] = {
-            'beats': hbt.onsets,
-            'color': next(colors)
-    }
-
-    for name, bt in beat_trackers.items():
-        print('{0}: {1}'.format(name, bt))
+    colors = [
+        (255, 0, 0),  # red
+        (0, 255, 0),  # green
+        (0, 165, 255),  # blue
+        (255, 0, 255),  # magenta
+        (255, 255, 0),  # yellow
+        (255, 69, 0),  # orange
+        (0, 255, 255),  # cyan
+        (145, 112, 235),  # blue-violet
+    ]
+    colorcycle = itertools.cycle(colors)
 
     fps = 30
 
@@ -76,34 +71,93 @@ def main():
     frame_duration = 1 / fps
     frame_duration_ms = frame_duration * 1000
 
-    total_duration = float(audio.shape[0])/44100.0
-    total_frames = total_duration/frame_duration
+    total_duration = numpy.floor(float(audio.shape[0]) / 44100.0)
+    total_frames = int(numpy.ceil(total_duration / frame_duration))
+
+    times_vector = numpy.arange(0, total_duration, frame_duration)
+
+    all_beat_times = individual_tracker_beat_locations + [
+        strong_beat_locations,
+        hbt.beat_consensus,
+    ]
+
+    all_beat_frames = [
+        numpy.concatenate(
+            (
+                numpy.zeros(
+                    1,
+                ),
+                find_closest(times_vector, beat_times),
+                numpy.ones(
+                    1,
+                )
+                * (total_frames - 1),
+            )
+        ).astype(numpy.int)
+        for beat_times in all_beat_times
+    ]
+
+    off_beat_frames = [
+        ((x[1:] + x[:-1]) / 2).astype(numpy.int) for x in all_beat_frames
+    ]
+
+    all_positions = []  # []
+    for i in range(len(all_beat_frames)):
+        x = (
+            numpy.empty(
+                total_frames,
+            )
+            * numpy.nan
+        )
 
-    total_duration = frame_duration * total_frames
+        x[all_beat_frames[i]] = 1
+        x[off_beat_frames[i]] = -1
+        a = pd.Series(x)
+        all_positions.append(a.interpolate().to_numpy())
 
     blank_frame = numpy.zeros((video_height, video_width, 3), numpy.uint8)
 
-    def render_animations(*args, **kwargs):
-        video_frame = blank_frame.copy()
+    box_width = int(video_width / 4)
+    box_edges_horiz = numpy.arange(0, video_width + 1, box_width)
+    box_centers_horiz = box_edges_horiz[:-1] + int(box_width / 2)
 
-        # draw stick figures with text
-
-        # draw some text, names of algorithms etc.
-        cv2.putText(
-            video_frame,
-            "BEAT",
-            all_beat_pos,
-            cv2.FONT_HERSHEY_SIMPLEX,
-            2.0,
-            all_beat_color,
-            3,
-            cv2.LINE_AA,
-        )
-        cv2.line(image, (20,10), (100,10), (255,0,0), 2)
+    box_height = int(video_height / 2)
+    box_edges_vert = numpy.arange(0, video_height + 1, box_height)
+    box_centers_vert = box_edges_vert[:-1] + int(box_height / 2)
+
+    positions = list(itertools.product(box_centers_horiz, box_centers_vert))
+
+    frame_index = 0
 
+    def render_animations(*args, **kwargs):
+        nonlocal frame_index
+        video_frame = blank_frame.copy()
 
-        # adjust color on frames
-        video_frame = cv2.cvtColor(video_frame, cv2.COLOR_BGR2RGB)
+        for i, beats in enumerate(all_beat_frames):
+            center = positions[i]
+            try:
+                interpolated_pos = all_positions[i][frame_index]
+            except IndexError:
+                interpolated_pos = 0
+
+            current_position = (
+                center[0],
+                int(center[1] + (box_height / 2 - 100) * interpolated_pos),
+            )
+
+            # draw some text, names of algorithms etc.
+            cv2.putText(
+                video_frame,
+                str(i),
+                current_position,
+                cv2.FONT_HERSHEY_SIMPLEX,
+                2.0,
+                colors[i],
+                3,
+                cv2.LINE_AA,
+            )
+
+        frame_index += 1
         return video_frame
 
     print("Processing video - rendering animations")

requirements.txt

@@ -11,3 +11,4 @@ moviepy==1.0.3
 essentia==2.1b6.dev374
 madmom==0.16.1
 scipy==1.6.0
+pandas