File size: 4,351 Bytes
72f552e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
"""Lyrics-to-beat mapping: group beats into segments and assign lyrics."""

import json
from pathlib import Path
from typing import Optional


def segment_lyrics(
    beats: list[dict],
    lyrics: list[dict],
    beats_per_segment: int = 4,
) -> list[dict]:
    """Map timestamped lyrics onto beat-grouped segments.

    Groups consecutive beats into segments (e.g. 4 beats = 1 bar in 4/4 time)
    and assigns words to the segment where they start.

    Args:
        beats: List of beat dicts with "beat" and "time" keys.
        lyrics: List of word dicts with "word", "start", "end" keys.
        beats_per_segment: Number of beats per segment. 4 = one bar in 4/4 time.

    Returns:
        List of segment dicts with keys:
            - segment: 1-indexed segment number
            - start: start time in seconds
            - end: end time in seconds
            - duration: segment duration in seconds
            - lyrics: raw lyrics text for this segment (may be empty)
            - words: list of word dicts that fall in this segment
    """
    beat_times = [b["time"] for b in beats]

    # Build segment boundaries by grouping every N beats
    segments = []
    seg_num = 1
    for i in range(0, len(beat_times) - 1, beats_per_segment):
        start = beat_times[i]
        # End is either N beats later or the last beat
        end_idx = min(i + beats_per_segment, len(beat_times) - 1)
        end = beat_times[end_idx]

        # Store individual beat timestamps for this segment
        seg_beat_times = [
            round(beat_times[j], 3)
            for j in range(i, min(i + beats_per_segment + 1, len(beat_times)))
        ]

        segments.append({
            "segment": seg_num,
            "start": round(start, 3),
            "end": round(end, 3),
            "duration": round(end - start, 3),
            "beats": seg_beat_times,
            "lyrics": "",
            "words": [],
        })
        seg_num += 1

    # Assign words to segments based on where the word starts
    for word in lyrics:
        word_start = word["start"]
        for seg in segments:
            if seg["start"] <= word_start < seg["end"]:
                seg["words"].append(word)
                break
        else:
            # Word starts after last segment boundary — assign to last segment
            if segments and word_start >= segments[-1]["start"]:
                segments[-1]["words"].append(word)

    # Build lyrics text per segment
    for seg in segments:
        seg["lyrics"] = " ".join(w["word"] for w in seg["words"])

    return segments


def save_segments(
    segments: list[dict],
    output_path: str | Path,
) -> Path:
    """Save segments to a JSON file.

    Args:
        segments: List of segment dicts.
        output_path: Path to save the JSON file.

    Returns:
        Path to the saved JSON file.
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    with open(output_path, "w") as f:
        json.dump(segments, f, indent=2)

    return output_path


def run(
    data_dir: str | Path,
    beats_per_segment: int = 4,
) -> list[dict]:
    """Full segmentation pipeline: load beats + lyrics, segment, and save.

    Args:
        data_dir: Song data directory containing beats.json and lyrics.json
            (e.g. data/Gone/).
        beats_per_segment: Number of beats per segment (4 = one bar).

    Returns:
        List of segment dicts.
    """
    data_dir = Path(data_dir)

    with open(data_dir / "beats.json") as f:
        beats = json.load(f)

    with open(data_dir / "lyrics.json") as f:
        lyrics = json.load(f)

    segments = segment_lyrics(beats, lyrics, beats_per_segment=beats_per_segment)
    save_segments(segments, data_dir / "segments.json")

    return segments


if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python -m src.segmenter <data_dir>")
        print("  e.g. python -m src.segmenter data/Gone")
        sys.exit(1)

    segments = run(sys.argv[1])
    print(f"Created {len(segments)} segments:\n")
    for seg in segments:
        lyrics_display = f'"{seg["lyrics"]}"' if seg["lyrics"] else "(instrumental)"
        print(f"  Seg {seg['segment']}: {seg['start']:.3f}s - {seg['end']:.3f}s "
              f"({seg['duration']:.3f}s) {lyrics_display}")