File size: 6,728 Bytes
d4004bc 6e1c8e6 d4004bc 6e1c8e6 d4004bc 6e1c8e6 d4004bc 6e1c8e6 d4004bc 6e1c8e6 d4004bc |
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 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 |
"""
Copyright $today.year LY Corporation
LY Corporation licenses this file to you under the Apache License,
version 2.0 (the "License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at:
https://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
License for the specific language governing permissions and limitations
under the License.
"""
import os
import torch
import subprocess
import gradio as gr
import librosa
from tqdm import tqdm
from lighthouse.models import *
# use GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
MODEL_NAMES = ['qd_detr']
FEATURES = ['clap']
TOPK_MOMENT = 5
sample_path = "sample_data/1a-ODBWMUAE.wav"
sample_query = "Water cascades down from a waterfall."
"""
Helper functions
"""
def load_pretrained_weights():
file_urls = []
for model_name in MODEL_NAMES:
for feature in FEATURES:
file_urls.append(
"https://zenodo.org/records/13961029/files/{}_{}_clotho-moment.ckpt".format(feature, model_name)
)
for file_url in tqdm(file_urls):
if not os.path.exists('gradio_demo/weights/' + os.path.basename(file_url)):
command = 'wget -P gradio_demo/weights/ {}'.format(file_url)
subprocess.run(command, shell=True)
return file_urls
def flatten(array2d):
list1d = []
for elem in array2d:
list1d += elem
return list1d
"""
Model initialization
"""
load_pretrained_weights()
model = QDDETRPredictor('gradio_demo/weights/clap_qd_detr_clotho-moment.ckpt', device=device, feature_name='clap')
loaded_audio = None
"""
Gradio functions
"""
def audio_upload(audio):
global loaded_audio
if audio is None:
loaded_audio = None
yield gr.update(value="Removed the audio", visible=True)
else:
yield gr.update(value="Processing the audio. Wait for a minute...", visible=True)
audio_feats = model.encode_audio(audio)
loaded_audio = audio_feats
yield gr.update(value="Finished audio processing!", visible=True)
def model_load(radio):
if radio is not None:
yield gr.update(value="Loading new model. Wait for a minute...", visible=True)
global model
feature, model_name = radio.split('+')
feature, model_name = feature.strip(), model_name.strip()
if model_name == 'qd_detr':
model_class = QDDETRPredictor
else:
raise gr.Error("Select from the models")
model = model_class('gradio_demo/weights/{}_{}_clotho-moment.ckpt'.format(feature, model_name),
device=device, feature_name='{}'.format(feature))
yield gr.update(value="Model loaded: {}".format(radio), visible=True)
def predict(textbox, line, gallery):
global loaded_audio
if loaded_audio is None:
raise gr.Error('Upload the audio before pushing the `Retrieve moment` button.')
else:
prediction = model.predict(textbox, loaded_audio)
mr_results = prediction['pred_relevant_windows']
buttons = []
for i, pred in enumerate(mr_results[:TOPK_MOMENT]):
buttons.append(gr.Button(value='moment {}: [{}, {}] Score: {}'.format(i+1, pred[0], pred[1], pred[2]), visible=True))
return buttons
def show_trimmed_audio(audio, button):
s, sr = librosa.load(audio, sr=None)
_seconds = button.split(': [')[1].split(']')[0].split(', ')
start_sec = float(_seconds[0])
end_sec = float(_seconds[1])
start_frame = int(start_sec * sr)
end_frame = int(end_sec * sr)
return gr.Audio((sr, s[start_frame:end_frame]), interactive=False, visible=True)
def main():
title = """# Audio Moment Retrieval Demo"""
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown(title)
with gr.Row():
with gr.Column():
with gr.Group():
gr.Markdown("## Model selection")
radio_list = flatten([["{} + {}".format(feature, model_name) for model_name in MODEL_NAMES] for feature in FEATURES])
radio = gr.Radio(radio_list, label="models", value="clap + qd_detr", info="Which model do you want to use?")
load_status_text = gr.Textbox(label='Model load status', value='Model loaded: clap + qd_detr')
with gr.Group():
gr.Markdown("## Audio and query")
audio_input = gr.Audio(sample_path, type='filepath')
output = gr.Textbox(label='Audio processing progress')
query_input = gr.Textbox(sample_query, label='query')
button = gr.Button("Retrieve moment", variant="primary")
with gr.Column():
with gr.Group():
gr.Markdown("## Retrieved moments")
gr.Markdown("Click on the moment button to listen to the trimmed audio.")
button_1 = gr.Button(value='moment 1', visible=False, elem_id='result_0')
button_2 = gr.Button(value='moment 2', visible=False, elem_id='result_1')
button_3 = gr.Button(value='moment 3', visible=False, elem_id='result_2')
button_4 = gr.Button(value='moment 4', visible=False, elem_id='result_3')
button_5 = gr.Button(value='moment 5', visible=False, elem_id='result_4')
result = gr.Audio(None, label='Trimmed audio', interactive=False, visible=False)
button_1.click(show_trimmed_audio, inputs=[audio_input, button_1], outputs=[result])
button_2.click(show_trimmed_audio, inputs=[audio_input, button_2], outputs=[result])
button_3.click(show_trimmed_audio, inputs=[audio_input, button_3], outputs=[result])
button_4.click(show_trimmed_audio, inputs=[audio_input, button_4], outputs=[result])
button_5.click(show_trimmed_audio, inputs=[audio_input, button_5], outputs=[result])
audio_input.change(audio_upload, inputs=[audio_input], outputs=output)
radio.select(model_load, inputs=[radio], outputs=load_status_text)
button.click(predict,
inputs=[query_input],
outputs=[button_1, button_2, button_3, button_4, button_5])
demo.load(audio_upload, inputs=[audio_input], outputs=output)
demo.launch()
if __name__ == "__main__":
main()
|