youfufan
58014dfd09
Switches inference to a transformers pipeline (+pyannote segments), adds overlapping-chunk processing and CSV timecode export for DaVinci. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
262 lines
9.9 KiB
Python
Executable File
262 lines
9.9 KiB
Python
Executable File
import customtkinter as ctk
|
||
import tkinter as tk
|
||
from tkinter import filedialog, messagebox
|
||
import os
|
||
import csv
|
||
from datetime import timedelta
|
||
|
||
import math
|
||
import librosa
|
||
import soundfile as sf
|
||
|
||
from transformers import pipeline
|
||
from pyannote.core import Segment
|
||
|
||
# ------------------------
|
||
# Time formatting function
|
||
# ------------------------
|
||
def seconds_to_timecode(total_seconds):
|
||
"""
|
||
Convert float seconds -> "HH:MM:SS.mmm" for DaVinci Resolve CSV markers.
|
||
"""
|
||
td = timedelta(seconds=total_seconds)
|
||
hours, remainder = divmod(td.seconds, 3600)
|
||
minutes, secs = divmod(remainder, 60)
|
||
ms = int(td.microseconds / 1000)
|
||
return f"{hours:02d}:{minutes:02d}:{secs:02d}.{ms:03d}"
|
||
|
||
def chunk_audio(audio_path, chunk_length_s=10.0, stride_s=5.0, sr=16000):
|
||
"""
|
||
Loads the entire audio with librosa, then yields overlapping chunks
|
||
(chunk_data, chunk_start, chunk_end, sr).
|
||
|
||
- chunk_length_s = length (seconds) of each chunk
|
||
- stride_s = how much we "step back" within each chunk for overlap
|
||
e.g. chunk_length=10, stride=5 => chunk #0 covers
|
||
0-10s, chunk #1 covers 5-15s, chunk #2 covers 10-20s, etc.
|
||
- sr = sampling rate to load the audio
|
||
"""
|
||
audio, sr = librosa.load(audio_path, sr=sr)
|
||
total_len_s = len(audio) / sr
|
||
|
||
# Step in seconds, between the start points of consecutive chunks
|
||
# e.g. for chunk=10, stride=5 => step=10 - 5=5
|
||
# means next chunk starts 5s after the previous chunk start
|
||
step = chunk_length_s - stride_s
|
||
|
||
# Edge case: if stride_s >= chunk_length_s, you won't have overlap
|
||
if step <= 0:
|
||
step = chunk_length_s # no overlap
|
||
|
||
# Figure out how many chunks needed so we don’t exceed total length
|
||
# (a bit of math to ensure we handle trailing audio < chunk_length_s)
|
||
num_chunks = math.ceil((total_len_s - chunk_length_s) / step) + 1
|
||
if num_chunks < 1:
|
||
num_chunks = 1
|
||
|
||
for i in range(num_chunks):
|
||
chunk_start = i * step
|
||
chunk_end = chunk_start + chunk_length_s
|
||
|
||
# If the chunk goes beyond total_len_s, clamp it
|
||
if chunk_end > total_len_s:
|
||
chunk_end = total_len_s
|
||
|
||
# Convert times to sample indexes
|
||
start_sample = int(chunk_start * sr)
|
||
end_sample = int(chunk_end * sr)
|
||
|
||
# If our chunk_start >= total_len_s, we can stop
|
||
if chunk_start >= total_len_s:
|
||
break
|
||
|
||
# Slice the waveform
|
||
chunk_data = audio[start_sample:end_sample]
|
||
|
||
yield chunk_data, chunk_start, chunk_end, sr
|
||
|
||
class MusicMarkerApp(ctk.CTk):
|
||
def __init__(self):
|
||
super().__init__()
|
||
|
||
self.title("Music Marker Generator")
|
||
self.geometry("500x350")
|
||
|
||
# Variables to hold file paths & threshold
|
||
self.audio_file_path = tk.StringVar(value="")
|
||
self.output_csv_path = tk.StringVar(value="")
|
||
self.merge_threshold_var = tk.DoubleVar(value=2.0) # default: merge segments within 2 seconds
|
||
self.pipeline = None # We'll load it once on demand
|
||
|
||
self._create_widgets()
|
||
|
||
def _create_widgets(self):
|
||
# 1) Frame for selecting audio file
|
||
file_frame = ctk.CTkFrame(self)
|
||
file_frame.pack(pady=10, padx=10, fill="x")
|
||
|
||
file_label = ctk.CTkLabel(file_frame, text="Audio File:")
|
||
file_label.pack(side="left", padx=5)
|
||
|
||
file_entry = ctk.CTkEntry(file_frame, textvariable=self.audio_file_path, width=250)
|
||
file_entry.pack(side="left", padx=5)
|
||
|
||
file_button = ctk.CTkButton(file_frame, text="Browse", command=self._browse_audio_file)
|
||
file_button.pack(side="left", padx=5)
|
||
|
||
# 2) Frame for output CSV
|
||
output_frame = ctk.CTkFrame(self)
|
||
output_frame.pack(pady=10, padx=10, fill="x")
|
||
|
||
output_label = ctk.CTkLabel(output_frame, text="Output CSV:")
|
||
output_label.pack(side="left", padx=5)
|
||
|
||
output_entry = ctk.CTkEntry(output_frame, textvariable=self.output_csv_path, width=250)
|
||
output_entry.pack(side="left", padx=5)
|
||
|
||
output_button = ctk.CTkButton(output_frame, text="Browse", command=self._browse_output_csv)
|
||
output_button.pack(side="left", padx=5)
|
||
|
||
# 3) Merge Threshold
|
||
threshold_frame = ctk.CTkFrame(self)
|
||
threshold_frame.pack(pady=10, padx=10, fill="x")
|
||
|
||
threshold_label = ctk.CTkLabel(threshold_frame, text="Merge Gap (sec):")
|
||
threshold_label.pack(side="left", padx=5)
|
||
|
||
threshold_entry = ctk.CTkEntry(threshold_frame, textvariable=self.merge_threshold_var, width=50)
|
||
threshold_entry.pack(side="left", padx=5)
|
||
|
||
# 4) Run button
|
||
run_button = ctk.CTkButton(self, text="Run Music Detection", command=self._run_detection)
|
||
run_button.pack(pady=10)
|
||
|
||
# 5) Status label
|
||
self.status_label = ctk.CTkLabel(self, text="", wraplength=400, justify="left")
|
||
self.status_label.pack(pady=5)
|
||
|
||
def _browse_audio_file(self):
|
||
file_path = filedialog.askopenfilename(
|
||
title="Select Audio File",
|
||
filetypes=[("Audio Files", "*.wav *.mp3 *.flac *.m4a *.aac *.ogg *.wma *.aif *.aiff")]
|
||
)
|
||
if file_path:
|
||
self.audio_file_path.set(file_path)
|
||
|
||
def _browse_output_csv(self):
|
||
file_path = filedialog.asksaveasfilename(
|
||
title="Select Output CSV",
|
||
defaultextension=".csv",
|
||
filetypes=[("CSV Files", "*.csv")]
|
||
)
|
||
if file_path:
|
||
self.output_csv_path.set(file_path)
|
||
|
||
def _run_detection(self):
|
||
audio_path = self.audio_file_path.get().strip()
|
||
output_csv = self.output_csv_path.get().strip()
|
||
merge_threshold = self.merge_threshold_var.get()
|
||
|
||
if not audio_path or not os.path.isfile(audio_path):
|
||
messagebox.showerror("Error", "Please select a valid audio file.")
|
||
return
|
||
if not output_csv:
|
||
messagebox.showerror("Error", "Please specify an output CSV file.")
|
||
return
|
||
|
||
self._set_status("Loading model, please wait...")
|
||
|
||
# Load pipeline if not loaded yet
|
||
if self.pipeline is None:
|
||
try:
|
||
# Standard audio-classification pipeline
|
||
self.pipeline = pipeline(
|
||
"audio-classification",
|
||
model="MarekCech/GenreVim-Music-Detection-DistilHuBERT"
|
||
)
|
||
except Exception as e:
|
||
messagebox.showerror("Model Error", f"Could not load the model:\n{e}")
|
||
return
|
||
|
||
self._set_status("Chunking audio and running music detection...")
|
||
|
||
# We'll collect "music" segments from each chunk
|
||
music_segments = []
|
||
|
||
# Manually chunk the audio & classify each chunk
|
||
try:
|
||
for chunk_data, chunk_start, chunk_end, sr in chunk_audio(
|
||
audio_path,
|
||
chunk_length_s=10.0,
|
||
stride_s=5.0,
|
||
sr=16000
|
||
):
|
||
# The pipeline expects a waveform plus sampling_rate, not a path
|
||
# So pass the chunk_data + sr
|
||
results = self.pipeline(chunk_data, sampling_rate=sr)
|
||
if not results:
|
||
continue
|
||
|
||
# best guess from the pipeline for this chunk
|
||
best_guess = max(results, key=lambda x: x["score"])
|
||
if best_guess["label"].lower() == "music" and best_guess["score"] > 0.5:
|
||
music_segments.append(Segment(chunk_start, chunk_end))
|
||
|
||
except Exception as e:
|
||
messagebox.showerror("Detection Error", f"Error running music detection:\n{e}")
|
||
return
|
||
|
||
# Sort segments by start time
|
||
music_segments.sort(key=lambda seg: seg.start)
|
||
|
||
# Merge close segments based on threshold
|
||
merged_segments = []
|
||
if not music_segments:
|
||
self._set_status("No music segments found. No CSV generated.")
|
||
return
|
||
else:
|
||
current_start = music_segments[0].start
|
||
current_end = music_segments[0].end
|
||
|
||
for seg in music_segments[1:]:
|
||
# If the new segment starts within X seconds of the old segment, merge them
|
||
if seg.start <= current_end + merge_threshold:
|
||
current_end = max(current_end, seg.end)
|
||
else:
|
||
merged_segments.append(Segment(current_start, current_end))
|
||
current_start = seg.start
|
||
current_end = seg.end
|
||
|
||
# finalize last
|
||
merged_segments.append(Segment(current_start, current_end))
|
||
|
||
# Write CSV for DaVinci Resolve
|
||
try:
|
||
with open(output_csv, "w", newline="", encoding="utf-8") as csvfile:
|
||
writer = csv.writer(csvfile)
|
||
# DaVinci Resolve CSV marker columns
|
||
writer.writerow(["Name", "Start", "End", "Color", "Marker Type"])
|
||
|
||
for i, seg in enumerate(merged_segments, start=1):
|
||
start_tc = seconds_to_timecode(seg.start)
|
||
end_tc = seconds_to_timecode(seg.end)
|
||
name = f"Music Segment #{i}"
|
||
color = "Green"
|
||
marker_type = "Comment"
|
||
writer.writerow([name, start_tc, end_tc, color, marker_type])
|
||
|
||
self._set_status(
|
||
f"Done! Found {len(merged_segments)} music segments.\n"
|
||
f"CSV saved to: {output_csv}"
|
||
)
|
||
except Exception as e:
|
||
messagebox.showerror("File Error", f"Could not write CSV:\n{e}")
|
||
|
||
def _set_status(self, msg):
|
||
self.status_label.configure(text=msg)
|
||
self.update_idletasks()
|
||
|
||
if __name__ == "__main__":
|
||
app = MusicMarkerApp()
|
||
app.mainloop()
|