kundali.py (18743B)
1 #!/usr/bin/env python3 2 """ 3 kundali: scans track metadata and builds a sequence around an anchor track 4 """ 5 6 from __future__ import annotations 7 8 import argparse 9 import json 10 import logging 11 import re 12 import shutil 13 import subprocess 14 import sys 15 import tempfile 16 from pathlib import Path 17 18 import mutagen 19 20 PITCH_CLASSES = { 21 "C": 0, 22 "G": 1, 23 "D": 2, 24 "A": 3, 25 "E": 4, 26 "B": 5, 27 "F#": 6, 28 "Gb": 6, 29 "C#": 7, 30 "Db": 7, 31 "G#": 8, 32 "Ab": 8, 33 "D#": 9, 34 "Eb": 9, 35 "A#": 10, 36 "Bb": 10, 37 "F": 11, 38 } 39 40 CAMLOT_MAJOR = { 41 0: 8, 42 1: 9, 43 2: 10, 44 3: 11, 45 4: 12, 46 5: 1, 47 6: 2, 48 7: 3, 49 8: 4, 50 9: 5, 51 10: 6, 52 11: 7, 53 } 54 55 BPM_TAGS = ("TBPM", "BPM", "FBPM") 56 KEY_TAGS = ("TKEY", "INITIALKEY", "KEY", "FKEY") 57 TITLE_TAGS = ("TITLE", "TIT2", "TIT1") 58 59 CURVES = {"linear": "tri", "qsin": "qsin", "hsin": "hsin", "exp": "exp", "log": "log"} 60 CODECS = {"wav": "pcm_f32le", "flac": "flac", "mp3": "libmp3lame", "ogg": "libvorbis"} 61 62 SAMPLE_RATE = 48000 63 LOGGER = logging.getLogger("kundali") 64 65 _KEY_RE = re.compile(r"^(?P<note>.*?)\s*(?P<mode>MAJOR|MINOR|M|m)?$", re.IGNORECASE) 66 67 68 def error(msg: str) -> None: 69 LOGGER.error(msg) 70 sys.exit(1) 71 72 73 def parse_key(key_str: str) -> tuple[int, bool]: 74 original = key_str.strip() 75 match = _KEY_RE.match(original) 76 if not match: 77 raise ValueError(f"Unknown key: {original!r}") 78 79 note = match.group("note").strip().title() 80 mode = match.group("mode") 81 major = True 82 if mode is not None: 83 major = mode.upper() == "MAJOR" or mode == "M" 84 85 pitch = PITCH_CLASSES.get(note) 86 if pitch is None: 87 raise ValueError(f"Unknown key: {original!r}") 88 return pitch, major 89 90 91 def camelot(pitch: int, major: bool) -> str: 92 return f"{CAMLOT_MAJOR[pitch]}{'A' if major else 'B'}" 93 94 95 def read_tag(tags: dict[str, object], keys: tuple[str, ...]) -> str | None: 96 for key in keys: 97 if key not in tags: 98 continue 99 raw = tags[key] 100 if isinstance(raw, list): 101 raw = raw[0] if raw else "" 102 value = str(raw).strip() 103 if value: 104 return value 105 return None 106 107 108 def run_ffmpeg(*args: str) -> None: 109 cmd = ["ffmpeg", "-y", "-hide_banner", "-loglevel", "error", *args] 110 try: 111 subprocess.run(cmd, check=True, capture_output=True, text=True) 112 except subprocess.CalledProcessError as exc: 113 stderr = exc.stderr.strip() if exc.stderr else "(no stderr)" 114 error(f"ffmpeg failed: {stderr}") 115 except FileNotFoundError: 116 error("ffmpeg not found in PATH") 117 118 119 def format_filter(sample_rate: int, channels: int) -> str: 120 layout = "stereo" if channels == 2 else f"{channels}c" 121 return ( 122 f"aformat=sample_fmts=fltp:sample_rates={sample_rate}:channel_layouts={layout}" 123 ) 124 125 126 def load_tracks(directory: Path, recursive: bool = False) -> list[dict]: 127 if not directory.is_dir(): 128 error(f"Not a directory: {directory}") 129 130 files = directory.rglob("*") if recursive else directory.iterdir() 131 tracks = [] 132 133 for entry in sorted(files): 134 if not entry.is_file(): 135 continue 136 try: 137 audio = mutagen.File(str(entry)) 138 except Exception: 139 continue 140 if audio is None: 141 continue 142 143 tags = {k.upper(): v for k, v in audio.items()} 144 145 bpm_raw = read_tag(tags, BPM_TAGS) 146 key_raw = read_tag(tags, KEY_TAGS) 147 if bpm_raw is None or key_raw is None: 148 continue 149 150 try: 151 bpm = float(bpm_raw) 152 except ValueError: 153 continue 154 if bpm <= 0: 155 continue 156 157 try: 158 pitch, major = parse_key(key_raw) 159 except ValueError as exc: 160 LOGGER.warning("Skipping %s: %s", entry, exc) 161 continue 162 163 duration = getattr(getattr(audio, "info", None), "length", 0.0) 164 title = read_tag(tags, TITLE_TAGS) or entry.stem 165 166 tracks.append( 167 { 168 "path": entry.resolve(), 169 "bpm": bpm, 170 "key": key_raw, 171 "major": major, 172 "pitch": pitch, 173 "duration": duration, 174 "title": title, 175 "rel_path": str(entry.relative_to(directory)), 176 } 177 ) 178 179 if not tracks: 180 error(f"No tracks with BPM+key found in {directory}") 181 182 return tracks 183 184 185 def key_distance(a: dict, b: dict) -> int: 186 diff = abs((a["pitch"] - b["pitch"]) % 12) 187 diff = min(diff, 12 - diff) 188 189 if a["major"] != b["major"] and diff == 3: 190 return 0 191 192 return diff 193 194 195 def bpm_distance(a: dict, b: dict) -> float: 196 return abs(a["bpm"] - b["bpm"]) / max(a["bpm"], b["bpm"]) 197 198 199 def score_transition(a: dict, b: dict, bpm_weight: float = 10.0) -> float: 200 return key_distance(a, b) * 2.0 + bpm_distance(a, b) * bpm_weight 201 202 203 def build_set( 204 tracks: list[dict], 205 *, 206 first_track: str | None, 207 start_idx: int, 208 target_duration: float, 209 play_duration: float, 210 fade_duration: float, 211 bpm_window: float, 212 ) -> list[dict]: 213 pool = list(tracks) 214 215 if first_track is not None: 216 query = first_track.lower() 217 matches = [ 218 i 219 for i, t in enumerate(pool) 220 if query in t["title"].lower() 221 or query in t["path"].name.lower() 222 or query in t["rel_path"].lower() 223 ] 224 if not matches: 225 error(f"No track matching {first_track!r} found") 226 first_idx = matches[0] 227 else: 228 first_idx = start_idx % len(pool) 229 230 first = pool.pop(first_idx) 231 anchor_bpm = first["bpm"] 232 233 pool = [t for t in pool if abs(t["bpm"] - anchor_bpm) <= bpm_window] 234 if not pool: 235 error( 236 f"No tracks within {bpm_window:.1f} BPM of anchor track " 237 f"({first['title']}: {anchor_bpm:.1f} BPM). " 238 f"Widen with --bpm-window." 239 ) 240 241 sequence = [first] 242 first_play = min(play_duration, first["duration"]) 243 first_fade_out = max(0.0, first_play - fade_duration) 244 segments = [ 245 { 246 "track": first, 247 "start": 0.0, 248 "end": first_play, 249 "fade_out_start": first_fade_out, 250 } 251 ] 252 elapsed = first_fade_out 253 254 while elapsed < target_duration and pool: 255 last = sequence[-1] 256 best = min(pool, key=lambda candidate: score_transition(last, candidate)) 257 pool.remove(best) 258 sequence.append(best) 259 260 play = min(play_duration, best["duration"]) 261 start = elapsed 262 end = start + play 263 fade_out = max(0.0, end - fade_duration) 264 segments.append( 265 { 266 "track": best, 267 "start": start, 268 "end": end, 269 "fade_out_start": fade_out, 270 } 271 ) 272 elapsed = fade_out 273 274 return segments 275 276 277 def build_atempo_chain(ratio: float) -> list[str]: 278 stages = [] 279 remaining = ratio 280 while remaining > 2.0: 281 stages.append("atempo=2.0") 282 remaining /= 2.0 283 while remaining < 0.5: 284 stages.append("atempo=0.5") 285 remaining /= 0.5 286 stages.append(f"atempo={remaining:.6f}") 287 return stages 288 289 290 def fade_filters( 291 duration: float, fade_in: float, fade_out: float, curve: str 292 ) -> list[str]: 293 c = CURVES[curve] 294 295 fade_in = max(0.0, min(fade_in, duration / 2)) 296 fade_out = max(0.0, min(fade_out, duration / 2)) 297 fade_out_start = max(0.0, duration - fade_out) 298 299 return [ 300 f"afade=t=in:ss=0:d={fade_in:.3f}:curve={c}", 301 f"afade=t=out:st={fade_out_start:.3f}:d={fade_out:.3f}:curve={c}", 302 ] 303 304 305 def extract_segment( 306 src: Path, out: Path, duration: float, sample_rate: int, channels: int 307 ) -> None: 308 run_ffmpeg( 309 "-t", 310 f"{duration:.3f}", 311 "-i", 312 str(src), 313 "-ar", 314 str(sample_rate), 315 "-ac", 316 str(channels), 317 "-c:a", 318 "pcm_f32le", 319 str(out), 320 ) 321 322 323 def stretch_with_rubberband(src: Path, out: Path, ratio: float) -> None: 324 try: 325 subprocess.run( 326 ["rubberband", "-q", "-T", f"{ratio:.6f}", str(src), str(out)], 327 check=True, 328 capture_output=True, 329 text=True, 330 ) 331 except subprocess.CalledProcessError as exc: 332 stderr = exc.stderr.strip() if exc.stderr else "(no stderr)" 333 error(f"rubberband failed: {stderr}") 334 335 336 def render_segment( 337 segment: dict, 338 target_bpm: float, 339 *, 340 curve: str, 341 stretch_engine: str, 342 output_path: Path, 343 sample_rate: int, 344 channels: int, 345 ) -> None: 346 track = segment["track"] 347 play_duration = segment["end"] - segment["start"] 348 ratio = target_bpm / track["bpm"] 349 stretched_duration = play_duration * ratio 350 fade = min(play_duration / 2.0, play_duration - segment["fade_out_start"]) * ratio 351 352 if stretch_engine == "rubberband" and shutil.which("rubberband"): 353 with tempfile.TemporaryDirectory(prefix="dj_rb_") as tmpdir: 354 trimmed = Path(tmpdir) / "trimmed.wav" 355 extract_segment( 356 track["path"], trimmed, play_duration, sample_rate, channels 357 ) 358 stretched = Path(tmpdir) / "stretched.wav" 359 stretch_with_rubberband(trimmed, stretched, ratio) 360 filters = fade_filters(stretched_duration, fade, fade, curve) 361 filters.append(format_filter(sample_rate, channels)) 362 run_ffmpeg( 363 "-i", 364 str(stretched), 365 "-af", 366 ",".join(filters), 367 "-c:a", 368 "pcm_f32le", 369 str(output_path), 370 ) 371 return 372 373 if stretch_engine == "rubberband": 374 LOGGER.warning( 375 "rubberband not found, falling back to atempo for %s", track["title"] 376 ) 377 378 filters = [f"atrim=start=0.0:duration={play_duration:.3f}"] 379 if ratio != 1.0: 380 filters.extend(build_atempo_chain(ratio)) 381 filters.extend(fade_filters(stretched_duration, fade, fade, curve)) 382 filters.append(format_filter(sample_rate, channels)) 383 384 run_ffmpeg( 385 "-t", 386 f"{play_duration:.3f}", 387 "-i", 388 str(track["path"]), 389 "-af", 390 ",".join(filters), 391 "-c:a", 392 "pcm_f32le", 393 str(output_path), 394 ) 395 396 397 def _mix_pair( 398 a_path: Path, 399 b_path: Path, 400 b_delay: float, 401 out_path: Path, 402 sample_rate: int, 403 channels: int, 404 ) -> None: 405 delay_ms = max(0, int(b_delay * 1000)) 406 filter_complex = ( 407 f"[1:a]adelay=delays={delay_ms}|{delay_ms}:all=1[delayed];" 408 f"[0:a][delayed]amix=inputs=2:duration=longest:normalize=0[sum];" 409 f"[sum]{format_filter(sample_rate, channels)}[out]" 410 ) 411 run_ffmpeg( 412 "-i", 413 str(a_path), 414 "-i", 415 str(b_path), 416 "-filter_complex", 417 filter_complex, 418 "-map", 419 "[out]", 420 "-c:a", 421 "pcm_f32le", 422 str(out_path), 423 ) 424 425 426 def mix_segments( 427 segments: list[tuple[Path, float]], 428 output_path: Path, 429 *, 430 format_ext: str, 431 sample_rate: int, 432 channels: int, 433 ) -> None: 434 if not segments: 435 error("No segments to mix") 436 437 codec = CODECS.get(format_ext, "pcm_f32le") 438 quality = [] 439 if codec == "libmp3lame": 440 quality = ["-q:a", "2"] 441 elif codec == "libvorbis": 442 quality = ["-q:a", "6"] 443 444 with tempfile.TemporaryDirectory(prefix="kundali_mix_") as tmpdir: 445 acc_path = segments[0][0] 446 for i in range(1, len(segments)): 447 seg_path, start = segments[i] 448 mix_path = Path(tmpdir) / f"mix_{i:03d}.wav" 449 _mix_pair(acc_path, seg_path, start, mix_path, sample_rate, channels) 450 acc_path = mix_path 451 452 run_ffmpeg( 453 "-i", 454 str(acc_path), 455 "-c:a", 456 codec, 457 "-ar", 458 str(sample_rate), 459 "-ac", 460 str(channels), 461 *quality, 462 str(output_path), 463 ) 464 465 466 def positive_float(value: str) -> float: 467 f = float(value) 468 if f <= 0: 469 raise argparse.ArgumentTypeError(f"{value!r} must be positive") 470 return f 471 472 473 def nonnegative_float(value: str) -> float: 474 f = float(value) 475 if f < 0: 476 raise argparse.ArgumentTypeError(f"{value!r} must be non-negative") 477 return f 478 479 480 def parse_args(argv: list[str] | None = None) -> argparse.Namespace: 481 parser = argparse.ArgumentParser( 482 description="Match songs based on their key and BPM and merge them into an audio file", 483 formatter_class=argparse.ArgumentDefaultsHelpFormatter, 484 ) 485 parser.add_argument( 486 "--dir", 487 type=Path, 488 default=Path("."), 489 help="Directory with analyzed audio files", 490 ) 491 parser.add_argument( 492 "--recursive", action="store_true", help="Walk into subdirectories" 493 ) 494 parser.add_argument( 495 "--out", 496 type=Path, 497 default=Path("dj_set.wav"), 498 help="Output file (wav/flac/mp3/ogg)", 499 ) 500 parser.add_argument( 501 "--duration", 502 type=positive_float, 503 default=1800.0, 504 help="Target set length in seconds", 505 ) 506 parser.add_argument( 507 "--start", 508 type=int, 509 default=0, 510 help="Index of first track (when --first-track is not given)", 511 ) 512 parser.add_argument( 513 "--first-track", 514 type=str, 515 default=None, 516 help="Title/filename substring of the anchor track", 517 ) 518 parser.add_argument( 519 "--target-bpm", 520 type=positive_float, 521 default=None, 522 help="Override the common BPM", 523 ) 524 parser.add_argument( 525 "--play-duration", 526 type=positive_float, 527 default=240.0, 528 help="Max seconds per track", 529 ) 530 parser.add_argument( 531 "--fade", 532 type=nonnegative_float, 533 default=16.0, 534 help="Crossfade duration in seconds", 535 ) 536 parser.add_argument( 537 "--bpm-window", 538 type=nonnegative_float, 539 default=5.0, 540 help="Max BPM distance from the anchor track", 541 ) 542 parser.add_argument( 543 "--fade-curve", choices=list(CURVES), default="qsin", help="Crossfade curve" 544 ) 545 parser.add_argument( 546 "--stretch-engine", 547 choices=("atempo", "rubberband"), 548 default="atempo", 549 help="Time-stretch engine", 550 ) 551 parser.add_argument( 552 "--sample-rate", type=int, default=SAMPLE_RATE, help="Output sample rate in Hz" 553 ) 554 parser.add_argument( 555 "--channels", type=int, default=2, help="Output channel count (1 or 2)" 556 ) 557 parser.add_argument( 558 "--save-plan", type=Path, default=None, help="Write JSON tracklist/plan" 559 ) 560 parser.add_argument( 561 "-v", "--verbose", action="store_true", help="Enable debug logging" 562 ) 563 564 if argv is None: 565 argv = sys.argv[1:] 566 if not argv: 567 parser.print_help(sys.stderr) 568 sys.exit(0) 569 570 return parser.parse_args(argv) 571 572 573 def configure_logging(verbose: bool) -> None: 574 logging.basicConfig( 575 level=logging.DEBUG if verbose else logging.INFO, 576 format="%(levelname)s: %(message)s", 577 ) 578 579 580 def main(argv: list[str] | None = None) -> int: 581 args = parse_args(argv) 582 configure_logging(args.verbose) 583 584 tracks = load_tracks(args.dir, recursive=args.recursive) 585 tracks = list({t["path"]: t for t in tracks}.values()) 586 587 segments = build_set( 588 tracks, 589 first_track=args.first_track, 590 start_idx=args.start, 591 target_duration=args.duration, 592 play_duration=args.play_duration, 593 fade_duration=args.fade, 594 bpm_window=args.bpm_window, 595 ) 596 597 target_bpm = ( 598 args.target_bpm if args.target_bpm is not None else segments[0]["track"]["bpm"] 599 ) 600 format_ext = args.out.suffix.lstrip(".").lower() or "wav" 601 602 LOGGER.info("Found %d unique track(s)", len(tracks)) 603 LOGGER.info("Anchor BPM: %.1f (window: ±%.1f)", target_bpm, args.bpm_window) 604 LOGGER.info("Sample rate: %d Hz", args.sample_rate) 605 LOGGER.info("Channels: %d", args.channels) 606 LOGGER.info("Fade curve: %s", args.fade_curve) 607 LOGGER.info("Stretch: %s", args.stretch_engine) 608 LOGGER.info( 609 "Planned %d track(s) in set, total length ~%.1fs", 610 len(segments), 611 segments[-1]["end"], 612 ) 613 614 print("\nSequence:") 615 for i, seg in enumerate(segments): 616 tr = seg["track"] 617 display = ( 618 f"{tr['title']} ({tr['rel_path']})" 619 if tr["title"] != tr["rel_path"] 620 else tr["title"] 621 ) 622 print( 623 f"[{i:02d}] {tr['key']:12s} ({camelot(tr['pitch'], tr['major']):4s}) " 624 f"{tr['bpm']:5.1f} -> {target_bpm:5.1f} {display}" 625 ) 626 print() 627 628 if args.save_plan: 629 plan = { 630 "target_bpm": target_bpm, 631 "sample_rate": args.sample_rate, 632 "channels": args.channels, 633 "fade_curve": args.fade_curve, 634 "stretch_engine": args.stretch_engine, 635 "segments": [ 636 { 637 "title": seg["track"]["title"], 638 "path": str(seg["track"]["path"]), 639 "rel_path": seg["track"]["rel_path"], 640 "key": seg["track"]["key"], 641 "camelot": camelot(seg["track"]["pitch"], seg["track"]["major"]), 642 "bpm": seg["track"]["bpm"], 643 "start": round(seg["start"], 3), 644 "end": round(seg["end"], 3), 645 "fade_out_start": round(seg["fade_out_start"], 3), 646 } 647 for seg in segments 648 ], 649 } 650 args.save_plan.write_text(json.dumps(plan, indent=2), encoding="utf-8") 651 LOGGER.info("Saved plan to %s", args.save_plan) 652 653 with tempfile.TemporaryDirectory(prefix="kundali_") as tmpdir: 654 rendered = [] 655 for i, seg in enumerate(segments): 656 seg_path = Path(tmpdir) / f"seg_{i:03d}.wav" 657 LOGGER.info("Rendering [%02d] %s ...", i, seg["track"]["title"]) 658 render_segment( 659 seg, 660 target_bpm, 661 curve=args.fade_curve, 662 stretch_engine=args.stretch_engine, 663 output_path=seg_path, 664 sample_rate=args.sample_rate, 665 channels=args.channels, 666 ) 667 rendered.append((seg_path, seg["start"])) 668 669 LOGGER.info("Mixing %d segments into %s ...", len(rendered), args.out) 670 mix_segments( 671 rendered, 672 args.out, 673 format_ext=format_ext, 674 sample_rate=args.sample_rate, 675 channels=args.channels, 676 ) 677 678 LOGGER.info("Done: %s", args.out) 679 return 0 680 681 682 if __name__ == "__main__": 683 sys.exit(main())