feat: gradio app

.gitignore

@@ -0,0 +1,5 @@
+.DS_Store
+__pycache__/
+*.pyc
+
+task/

app.py

@@ -0,0 +1,334 @@
+from matplotlib import pyplot as plt
+from accelerate import Accelerator
+from zero import zero
+import gradio as gr
+from typing import Tuple
+import os
+from os import path
+from utils import plot_spec
+import librosa
+from hashlib import md5
+from demucs.separate import main as demucs
+from pyannote.audio import Pipeline
+from json import dumps, loads
+import shutil
+
+accelerator = Accelerator()
+device = accelerator.device
+print(f"Running on {device}")
+
+pipeline = Pipeline.from_pretrained(
+    "pyannote/speaker-diarization-3.1", use_auth_token=os.environ["HF_TOKEN"]
+)
+pipeline.to(device)
+
+
+tasks = []
+os.makedirs("task", exist_ok=True)
+for task in os.listdir("task"):
+    if path.isdir(path.join("task", task)):
+        tasks.append(task)
+
+
+def gen_task_id(location: str):
+    # use md5 hash of video file as task id
+    video = open(location, "rb").read()
+    return md5(video).hexdigest()
+
+
+def extract_audio(video: str) -> Tuple[str, str, str]:
+    task_id = gen_task_id(video)
+    os.makedirs(path.join("task", task_id), exist_ok=True)
+
+    videodest = path.join("task", task_id, "video.mp4")
+    if not path.exists(videodest):
+        shutil.copy(video, videodest)
+
+    wav48k = path.join("task", task_id, "extracted_48k.wav")
+    if not path.exists(wav48k):
+        os.system(
+            f"ffmpeg -i {videodest} -vn -ar 48000 task/{task_id}/extracted_48k.wav"
+        )
+
+    spec = path.join("task", task_id, "extracted_48k.png")
+    if not path.exists(spec):
+        y, sr = librosa.load(wav48k, sr=16000)
+        fig = plot_spec(y, sr)
+        fig.savefig(path.join("task", task_id, "extracted_48k.png"))
+        plt.close(fig)
+
+    return (task_id, wav48k, spec)
+
+
+@zero()
+def extract_vocals(task_id: str) -> Tuple[str, str]:
+    audio = path.join("task", task_id, "extracted_48k.wav")
+    if not path.exists(audio):
+        raise gr.Error("Audio file not found")
+
+    vocals = path.join("task", task_id, "htdemucs", "extracted_48k", "vocals.wav")
+
+    if not path.exists(vocals):
+        demucs(
+            [
+                "-d",
+                str(device),
+                "-n",
+                "htdemucs",
+                "--two-stems",
+                "vocals",
+                "-o",
+                path.join("task", task_id),
+                audio,
+            ]
+        )
+
+    spec = path.join("task", task_id, "vocals.png")
+    if not path.exists(spec):
+        y, sr = librosa.load(vocals, sr=16000)
+        fig = plot_spec(y, sr)
+        fig.savefig(path.join("task", task_id, "vocals.png"))
+        plt.close(fig)
+
+    return (vocals, spec)
+
+
+@zero(duration=60 * 2)
+def diarize_audio(task_id: str):
+    vocals = path.join("task", task_id, "htdemucs", "extracted_48k", "vocals.wav")
+    if not path.exists(vocals):
+        raise gr.Error("Vocals file not found")
+
+    diarization_json = path.join("task", task_id, "diarization.json")
+    if not path.exists(diarization_json):
+        result = pipeline(vocals)
+        with open(diarization_json, "w") as f:
+            diarization = []
+            for turn, _, speaker in result.itertracks(yield_label=True):
+                diarization.append(
+                    {
+                        "speaker": speaker,
+                        "start": turn.start,
+                        "end": turn.end,
+                        "duration": turn.duration,
+                    }
+                )
+            f.write(dumps(diarization))
+    with open(diarization_json, "r") as f:
+        diarization = loads(f.read())
+
+    filtered_json = path.join("task", task_id, "filtered_diarization.json")
+    if not path.exists(filtered_json):
+        # Filter out segments shorter than 2 second and group by speaker
+        filtered_segments = {}
+        for turn in diarization:
+            speaker = turn["speaker"]
+            if turn["duration"] >= 2.0:
+                if speaker not in filtered_segments:
+                    filtered_segments[speaker] = []
+                filtered_segments[speaker].append(turn)
+
+        # Filter out speakers with less than 60 seconds of speech
+        filtered_segments = {
+            speaker: segments
+            for speaker, segments in filtered_segments.items()
+            if sum(segment["duration"] for segment in segments) >= 60
+        }
+
+        with open(filtered_json, "w") as f:
+            f.write(dumps(filtered_segments))
+    with open(filtered_json, "r") as f:
+        filtered_segments = loads(f.read())
+
+    return filtered_segments
+
+
+def generate_clips(task_id: str, speaker: str) -> Tuple[str, str]:
+    video = path.join("task", task_id, "video.mp4")
+    if not path.exists(video):
+        raise gr.Error("Video file not found")
+
+    filtered_json = path.join("task", task_id, "filtered_diarization.json")
+    if not path.exists(filtered_json):
+        raise gr.Error("Diarization not found")
+
+    with open(filtered_json, "r") as f:
+        filtered_segments = loads(f.read())
+
+    if speaker not in filtered_segments:
+        raise gr.Error("Speaker not found")
+
+    mp4 = path.join("task", task_id, f"{speaker}.mp4")
+    if not path.exists(mp4):
+        cmd = f'ffmpeg -i {video} -filter_complex "'
+        for i, segment in enumerate(filtered_segments[speaker]):
+            start = segment["start"]
+            end = segment["end"]
+            cmd += f"[0:v]trim=start={start}:end={end},setpts=PTS-STARTPTS[v{i}];"
+            cmd += f"[0:a]atrim=start={start}:end={end},asetpts=PTS-STARTPTS[a{i}];"
+        for i in range(len(filtered_segments[speaker])):
+            cmd += f"[v{i}][a{i}]"
+        cmd += f'concat=n={len(filtered_segments[speaker])}:v=1:a=1[outv][outa]" -map [outv] -map [outa] -y {mp4}'
+        os.system(cmd)
+
+    segments = path.join("task", task_id, f"{speaker}")
+    if not path.exists(segments):
+        os.makedirs(segments)
+        for i, segment in enumerate(filtered_segments[speaker]):
+            start = segment["start"]
+            end = segment["end"]
+            name = path.join(segments, f"{i}_{start:.2f}_{end:.2f}.wav")
+            cmd = f"ffmpeg -i {video} -ss {start} -to {end} -f wav {name}"
+            os.system(cmd)
+
+    segments_zip = path.join("task", task_id, f"{speaker}.zip")
+    if not path.exists(segments_zip):
+        os.system(f"zip -r {segments_zip} {segments}")
+
+    return mp4, segments_zip
+
+
+with gr.Blocks() as app:
+    gr.Markdown("# Video Speaker Diarization")
+
+    gr.Markdown(
+        """
+    First, upload a video file. And let us do some inspection on the audio of the video.
+    """
+    )
+    original_video = gr.Video(label="Upload a video", show_download_button=True)
+    preprocess_btn = gr.Button(value="Pre Process", variant="primary")
+    preprocess_btn_label = gr.Markdown("Press the button!")
+
+    with gr.Column(visible=False) as preprocess_output:
+        gr.Markdown(
+            """
+        Now you can see the spectrogram of the extracted audio.
+
+        Next, let's remove the background music from the audio.
+        """
+        )
+        task_id = gr.Textbox(label="Task ID", visible=False)
+        extracted_audio = gr.Audio(label="Extracted Audio", type="filepath")
+        extracted_audio_spec = gr.Image(label="Extracted Audio Spectrogram")
+
+        extract_vocals_btn = gr.Button(
+            value="Remove Background Music", variant="primary"
+        )
+        extract_vocals_btn_label = gr.Markdown("Press the button!")
+
+    with gr.Column(visible=False) as extract_vocals_output:
+        vocals = gr.Audio(label="Vocals", type="filepath")
+        vocals_spec = gr.Image(label="Vocals Spectrogram")
+
+        diarize_btn = gr.Button(value="Diarize", variant="primary")
+        diarize_btn_label = gr.Markdown("Press the button!")
+
+    with gr.Column(visible=False) as diarize_output:
+        gr.Markdown(
+            """
+        Now you can select the speaker from the dropdown below to generate the clips of the speaker.
+        """
+        )
+        speaker_select = gr.Dropdown(label="Speaker", choices=[])
+        diarization_result = gr.Markdown("")
+
+        generate_clips_btn = gr.Button(value="Generate Clips", variant="primary")
+        generate_clips_btn_label = gr.Markdown("Press the button!")
+
+    with gr.Column(visible=False) as generate_clips_output:
+        speaker_clip = gr.Video(label="Speaker Clip")
+        speaker_clip_zip = gr.File(label="Download Audio Segments")
+
+    def preprocess(video: str):
+        task_id_val, extracted_audio_val, extracted_audio_spec_val = extract_audio(
+            video
+        )
+        return {
+            preprocess_output: gr.Column(visible=True),
+            task_id: task_id_val,
+            extracted_audio: extracted_audio_val,
+            extracted_audio_spec: extracted_audio_spec_val,
+            preprocess_btn_label: gr.Markdown("", visible=False),
+        }
+
+    preprocess_btn.click(
+        fn=preprocess,
+        inputs=[original_video],
+        outputs=[
+            preprocess_output,
+            task_id,
+            extracted_audio,
+            extracted_audio_spec,
+            preprocess_btn_label,
+        ],
+        api_name="preprocess",
+    )
+
+    def extract_vocals_fn(task_id: str):
+        vocals_val, vocals_spec_val = extract_vocals(task_id)
+        return {
+            extract_vocals_output: gr.Column(visible=True),
+            vocals: vocals_val,
+            vocals_spec: vocals_spec_val,
+            extract_vocals_btn_label: gr.Markdown("", visible=False),
+        }
+
+    extract_vocals_btn.click(
+        fn=extract_vocals_fn,
+        inputs=[task_id],
+        outputs=[extract_vocals_output, vocals, vocals_spec, extract_vocals_btn_label],
+        api_name="extract_vocals",
+    )
+
+    def diarize_fn(task_id: str):
+        filtered_segments = diarize_audio(task_id)
+        choices = []
+        for speaker in filtered_segments:
+            total = sum(segment["duration"] for segment in filtered_segments[speaker])
+            choices.append((f"{speaker} ({total:.2f}s)", speaker))
+
+        info = ""
+        for speaker, segments in filtered_segments.items():
+            total = sum(segment["duration"] for segment in segments)
+            info += f"### Speaker {speaker}: ({total:.2f}s)\n"
+            for segment in segments:
+                start = segment["start"]
+                end = segment["end"]
+                info += f"- {start:.2f} - {end:.2f} ({segment['duration']:.2f}s)\n"
+        return {
+            diarize_output: gr.Column(visible=True),
+            speaker_select: gr.Dropdown(label="Speaker", choices=choices),
+            diarization_result: gr.Markdown(info),
+            diarize_btn_label: gr.Markdown("", visible=False),
+        }
+
+    diarize_btn.click(
+        fn=diarize_fn,
+        inputs=[task_id],
+        outputs=[diarize_output, speaker_select, diarization_result, diarize_btn_label],
+        api_name="diarize",
+    )
+
+    def generate_clips_fn(task_id: str, speaker: str):
+        speaker_clip_val, zip_val = generate_clips(task_id, speaker)
+        return {
+            generate_clips_output: gr.Column(visible=True),
+            speaker_clip: speaker_clip_val,
+            speaker_clip_zip: zip_val,
+            generate_clips_btn_label: gr.Markdown("", visible=False),
+        }
+
+    generate_clips_btn.click(
+        fn=generate_clips_fn,
+        inputs=[task_id, speaker_select],
+        outputs=[
+            generate_clips_output,
+            speaker_clip,
+            speaker_clip_zip,
+            generate_clips_btn_label,
+        ],
+        api_name="generate_clips",
+    )
+
+    app.launch()

requirements.txt

@@ -0,0 +1,10 @@
+torch==2.2.0
+soundfile==0.12.1
+numpy==1.26.0
+librosa==0.9.2
+einops==0.8.0
+gradio==4.37.2
+accelerate==0.31.0
+matplotlib==3.8.3
+demucs==4.0.1
+pyannote-audio==3.3.1

utils.py

@@ -0,0 +1,18 @@
+import numpy as np
+import librosa
+import matplotlib.pyplot as plt
+from librosa.display import specshow
+
+
+def plot_spec(y: np.ndarray, sr: int, title: str = "Spectrogram") -> plt.Figure:
+    y[np.isnan(y)] = 0
+    y[np.isinf(y)] = 0
+    stft = librosa.stft(y=y)
+    D = librosa.amplitude_to_db(np.abs(stft), ref=np.max)
+
+    fig = plt.figure(figsize=(10, 4))
+    specshow(D, sr=sr, y_axis="linear", x_axis="time", cmap="viridis")
+    plt.title(title)
+    plt.tight_layout()
+
+    return fig

zero.py

@@ -0,0 +1,21 @@
+import os
+
+zero_is_available = "SPACES_ZERO_GPU" in os.environ
+
+if zero_is_available:
+    import spaces  # type: ignore
+
+    print("ZeroGPU is available")
+else:
+    print("ZeroGPU is not available")
+
+
+# a decorator that applies the spaces.GPU decorator if zero is available
+def zero(duration=60):
+    def wrapper(func):
+        if zero_is_available:
+            return spaces.GPU(func, duration=duration)
+        else:
+            return func
+
+    return wrapper