| """ |
| 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 * |
|
|
| |
| 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() |
|
|
