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