Initial commit

Author: nik-418

Dockerfile

@@ -0,0 +1,27 @@
+# Must use a Cuda version 11+
+FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime
+
+ARG HF_TOKEN
+
+WORKDIR /
+
+# Install git and wget
+RUN apt-get update && apt-get install -y git wget
+
+# Upgrade pip
+RUN pip install --upgrade pip
+
+ADD requirements.txt requirements.txt
+RUN pip3 install -r requirements.txt
+
+ENV HF_TOKEN=${HF_TOKEN}
+
+# Add your model weight files 
+ADD download.py .
+RUN python3 download.py
+
+ADD . .
+
+EXPOSE 8000
+
+CMD python3 -u app.py

README.md

@@ -0,0 +1,39 @@
+![](https://www.banana.dev/lib_zOkYpJoyYVcAamDf/x2p804nk9qvjb1vg.svg?w=340 "Banana.dev")
+
+# Banana.dev bark starter template
+
+This is a segmentation starter template from [Banana.dev](https://www.banana.dev) that allows on-demand serverless GPU inference.
+
+You can fork this repository and deploy it on Banana as is, or customize it based on your own needs.
+
+
+# Running this app
+
+## Deploying on Banana.dev
+
+1. [Fork this](https://github.com/bananaml/demo-segmentation/fork) repository to your own GitHub account.
+2. Connect your GitHub account on Banana.
+3. [Create a new model](https://app.banana.dev/deploy) on Banana from the forked GitHub repository.
+
+## Running after deploying
+
+1. Wait for the model to build after creating it.
+2. Make an API request using one of the provided snippets in your Banana dashboard. However, instead of sending a prompt as provided in the snippet, adjust the prompt to fit the needs of the segmentation model:
+
+```python
+inputs = {
+    "audio": "bucket_link_to_wav_file",
+    "option": "voice_activity_detection"
+}
+```
+
+The `audio` parameter should be substituted with your S3 (or any other provider where you can store .wav files) bucket link that contains the .wav audio file you want to segment. For the `option` parameter, you have to choose between the following options depending on what segmentation information you want to gain from the audio file:
+
+* voice_activity_detection
+* overlapped_speech_detection
+* instantaneous_speaker_counting
+* speaker_change_detection
+
+In the example above, we chose `voice_activity_detection` as an option.
+
+For more info, check out the [Banana.dev docs](https://docs.banana.dev/banana-docs/).

app.py

@@ -0,0 +1,89 @@
+from potassium import Potassium, Request, Response
+from pyannote.audio import Model, Inference
+import numpy as np
+import os
+import base64
+import matplotlib.pyplot as plt
+import io
+import requests
+from pyannote.core import notebook
+
+MODEL = "pyannote/segmentation"
+HF_TOKEN = os.getenv('HF_TOKEN')
+SPEAKER_AXIS = 2
+TIME_AXIS = 1
+
+app = Potassium("segmentation")
+
+@app.init
+def init():
+    """
+    Initialize the application with the pretrained model.
+    """
+    model = Model.from_pretrained(MODEL, use_auth_token=HF_TOKEN)
+    context = {
+        "model": model
+    }
+    return context
+
+def process_audio(model, hook, label, filename):
+    """
+    Process the audio file with the given model, hook, label, and filename.
+    """
+    inference = Inference(model, pre_aggregation_hook=hook)
+    prob = inference(filename)
+    prob.labels = [label]
+    fig, ax = plt.subplots()
+    notebook.plot_feature(prob, ax=ax)
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    plt.close(fig)
+    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
+    return image_base64
+
+def voice_activity_detection(model):
+    to_vad = lambda o: np.max(o, axis=SPEAKER_AXIS, keepdims=True)
+    result = process_audio(model, to_vad, 'SPEECH', "/tmp/temp.wav")
+    return result
+
+def overlapped_speech_detection(model):
+    to_osd = lambda o: np.partition(o, -2, axis=SPEAKER_AXIS)[:, :, -2, np.newaxis]
+    result = process_audio(model, to_osd, 'OVERLAP', "/tmp/temp.wav")
+    return result
+
+def instantaneous_speaker_counting(model):
+    to_cnt = lambda probability: np.sum(probability, axis=SPEAKER_AXIS, keepdims=True)
+    result = process_audio(model, to_cnt, 'SPEAKER_COUNT', "/tmp/temp.wav")
+    return result
+
+def speaker_change_detection(model):
+    to_scd = lambda probability: np.max(
+        np.abs(np.diff(probability, n=1, axis=TIME_AXIS)), 
+        axis=SPEAKER_AXIS, keepdims=True)
+    result = process_audio(model, to_scd, 'SPEAKER_CHANGE', "/tmp/temp.wav")
+    return result
+
+@app.handler()
+def handler(context: dict, request: Request) -> Response:
+    """
+    Handle the incoming request and return the response.
+    """
+    model = context.get("model")
+    audio_input = request.json.get("audio")
+    option = request.json.get("option")
+    response = requests.get(audio_input)
+    filename = '/tmp/temp.wav'
+    with open(filename, 'wb') as f:
+        f.write(response.content)
+    if option == "voice_activity_detection":
+        image_base64 = voice_activity_detection(model)
+    elif option == "overlapped_speech_detection":
+        image_base64 = overlapped_speech_detection(model)
+    elif option == "instantaneous_speaker_counting":
+        image_base64 = instantaneous_speaker_counting(model)
+    elif option == "speaker_change_detection":
+        image_base64 = speaker_change_detection(model)
+    return Response(json={"output": image_base64}, status=200)  
+
+if __name__ == "__main__":
+    app.serve()

banana_config.json

@@ -0,0 +1,27 @@
+{
+    "version": "1",
+    "inputs": [
+        {
+            "name": "option",
+            "description": "Choose a segmentation to perform on the audio file",
+            "type": "options",
+            "options": ["voice_activity_detection","overlapped_speech_detection", "instantaneous_speaker_counting", "speaker_change_detection"],
+            "required": true,
+            "default": "voice_activity_detection"
+        },
+        {
+            "name": "audio",
+            "type": "file",
+            "mimes": ["audio/wav"],
+            "required": true
+        }
+    ],
+    "output": {
+        "image": {
+            "type": "image",
+            "source": "base64",
+            "path": "output"
+        }
+    }
+}
+

download.py

@@ -0,0 +1,11 @@
+from pyannote.audio import Model
+import os
+
+HF_TOKEN = os.getenv('HF_TOKEN')
+
+def download_model():
+    """Load model"""
+    model = Model.from_pretrained("pyannote/segmentation", use_auth_token=HF_TOKEN)
+
+if __name__ == "__main__":
+    download_model()

requirements.txt

@@ -0,0 +1,4 @@
+potassium>=0.0.9
+torch
+transformers
+https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip