initial

.gitignore

@@ -0,0 +1,4 @@
+.DS_Store
+__pycache__
+
+

README.md

@@ -1,12 +1,24 @@
 ---
+
 title: Sound Effect Search
+
 emoji: 🏢
+
 colorFrom: gray
+
 colorTo: green
+
+python: 3.9.7
+
 sdk: gradio
-sdk_version: 3.28.0
+
+sdk_version: 3.23.0
+
 app_file: app.py
-pinned: false
+
+pinned: true
+
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+

app.py

@@ -0,0 +1,40 @@
+import gradio as gr
+from src.search import Search
+
+search = Search("config.yaml")
+
+with gr.Blocks() as demo:
+    gr.Markdown("Search Sound Effect using this demo.")
+    with gr.TabItem("Search from Audio File"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                text_input = gr.Textbox(value="太鼓", label="SE Title")
+                audio_input = gr.Audio(source="upload")
+                ratio = gr.Slider(minimum=0, maximum=1, value=1, label="Weight Parameter. 1 means 'use only text'. 0 means 'use only audio'.")
+                topk = gr.Dropdown(
+                    [5, 10, 20, 30, 40, 50], value="20", label="Top K"
+                )
+                button = gr.Button("Search")
+            with gr.Column(scale=2):
+                output = gr.Dataframe()
+    with gr.TabItem("Search from Microphone"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                mic_text_input = gr.Textbox(value="太鼓", label="SE Title")
+                mic_audio_input = gr.Audio(source="microphone")
+                mic_ratio = gr.Slider(minimum=0, maximum=1, value=1, label="Weight Parameter. 1 means 'use only text'. 0 means 'use only audio'.")
+                mic_topk = gr.Dropdown(
+                    [5, 10, 20, 30, 40, 50], value="20", label="Top K"
+                )
+                mic_button = gr.Button("Search")
+            with gr.Column(scale=2):
+                mic_output = gr.Dataframe()
+
+    button.click(
+        search.search, inputs=[text_input, audio_input, ratio, topk], outputs=output
+    )
+    mic_button.click(
+        search.search, inputs=[mic_text_input, mic_audio_input, mic_ratio, mic_topk], outputs=mic_output
+    )
+
+demo.launch()

config.yaml

@@ -0,0 +1,5 @@
+path_data: data
+path_csv: result.csv
+path_text_embedding: text.pt
+path_audio_embedding: audio.pt
+sample_rate: 16000

requirements.txt

@@ -0,0 +1,82 @@
+aiofiles==23.1.0
+aiohttp==3.8.4
+aiosignal==1.3.1
+altair==4.2.2
+antlr4-python3-runtime==4.9.3
+anyio==3.6.2
+async-timeout==4.0.2
+attrs==23.1.0
+certifi==2022.12.7
+cffi==1.15.1
+charset-normalizer==3.1.0
+click==8.1.3
+contourpy==1.0.7
+cycler==0.11.0
+entrypoints==0.4
+fastapi==0.95.1
+ffmpy==0.3.0
+filelock==3.12.0
+fonttools==4.39.3
+frozenlist==1.3.3
+fsspec==2023.4.0
+fugashi==1.2.1
+gradio==3.28.1
+gradio_client==0.1.4
+h11==0.14.0
+httpcore==0.17.0
+httpx==0.24.0
+huggingface-hub==0.14.1
+idna==3.4
+importlib-resources==5.12.0
+ipadic==1.0.0
+Jinja2==3.1.2
+jsonschema==4.17.3
+kiwisolver==1.4.4
+linkify-it-py==2.0.0
+markdown-it-py==2.2.0
+MarkupSafe==2.1.2
+matplotlib==3.7.1
+mdit-py-plugins==0.3.3
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.0.4
+networkx==3.1
+numpy==1.24.3
+omegaconf==2.3.0
+orjson==3.8.11
+packaging==23.1
+pandas==2.0.1
+Pillow==9.5.0
+pycparser==2.21
+pydantic==1.10.7
+pydub==0.25.1
+pyparsing==3.0.9
+pyrsistent==0.19.3
+PySoundFile==0.9.0.post1
+python-dateutil==2.8.2
+python-multipart==0.0.6
+pytz==2023.3
+PyYAML==6.0
+regex==2023.3.23
+requests==2.29.0
+semantic-version==2.10.0
+six==1.16.0
+sniffio==1.3.0
+starlette==0.26.1
+sympy==1.11.1
+tokenizers==0.13.3
+toolz==0.12.0
+torch==2.0.0
+torchaudio==2.0.1
+torchvision==0.15.1
+tqdm==4.65.0
+transformers==4.28.1
+typing_extensions==4.5.0
+tzdata==2023.3
+uc-micro-py==1.0.1
+unidic-lite==1.0.8
+urllib3==1.26.15
+uvicorn==0.22.0
+websockets==11.0.2
+yarl==1.9.2
+zipp==3.15.0

src/create_embed.py

@@ -0,0 +1,125 @@
+import argparse
+import os
+
+import numpy as np
+import pandas as pd
+import torch
+from omegaconf import OmegaConf
+from pydub import AudioSegment
+from tqdm import trange
+from transformers import (
+    AutoFeatureExtractor,
+    BertForSequenceClassification,
+    BertJapaneseTokenizer,
+    Wav2Vec2ForXVector,
+)
+
+
+class Embeder:
+    def __init__(self, config):
+        self.config = OmegaConf.load(config)
+        self.df = pd.read_csv(config.path_csv)
+        self.audio_feature_extractor = AutoFeatureExtractor.from_pretrained(
+            "anton-l/wav2vec2-base-superb-sv"
+        )
+        self.audio_model = Wav2Vec2ForXVector.from_pretrained(
+            "anton-l/wav2vec2-base-superb-sv"
+        )
+        self.text_tokenizer = BertJapaneseTokenizer.from_pretrained(
+            "cl-tohoku/bert-base-japanese-whole-word-masking"
+        )
+        self.text_model = BertForSequenceClassification.from_pretrained(
+            "cl-tohoku/bert-base-japanese-whole-word-masking",
+            num_labels=2,
+            output_attentions=False,
+            output_hidden_states=True,
+        ).eval()
+
+    def run(self):
+        self._create_audio_embed()
+        self._create_text_embed()
+
+    def _create_audio_embed(self):
+        audio_embed = None
+        idx = []
+        for i in trange(len(self.df)):
+            audio = []
+            song = AudioSegment.from_wav(
+                os.path.join(
+                    self.config.path_data,
+                    "new_" + self.df.iloc[i]["filename"].replace(".mp3", ".wav"),
+                )
+            )
+            song = np.array(song.get_array_of_samples(), dtype="float")
+            audio.append(song)
+            inputs = self.audio_feature_extractor(
+                audio,
+                sampling_rate=self.config.sample_rate,
+                return_tensors="pt",
+                padding=True,
+            )
+            try:
+                with torch.no_grad():
+                    embeddings = self.audio_model(**inputs).embeddings
+                audio_embed = (
+                    embeddings
+                    if audio_embed is None
+                    else torch.concatenate([audio_embed, embeddings])
+                )
+            except Exception:
+                idx.append(i)
+
+        audio_embed = torch.nn.functional.normalize(audio_embed, dim=-1).cpu()
+        self.clean_and_save_data(audio_embed, idx)
+        self.df = self.df.drop(index=idx)
+        self.df.to_csv(self.config.path_csv, index=False)
+
+    def _create_text_embed(self):
+        text_embed = None
+        for i in range(len(self.df)):
+            sentence = self.df.iloc[i]["filename"].replace(".mp3", "")
+            tokenized_text = self.text_tokenizer.tokenize(sentence)
+            indexed_tokens = self.text_tokenizer.convert_tokens_to_ids(tokenized_text)
+            tokens_tensor = torch.tensor([indexed_tokens])
+            with torch.no_grad():
+                all_encoder_layers = self.text_model(tokens_tensor)
+            embedding = torch.mean(all_encoder_layers[1][-2][0], axis=0).reshape(1, -1)
+            text_embed = (
+                embedding
+                if text_embed is None
+                else torch.concatenate([text_embed, embedding])
+            )
+        text_embed = torch.nn.functional.normalize(text_embed, dim=-1).cpu()
+        torch.save(text_embed, self.config.path_text_embedding)
+
+    def clean_and_save_data(self, audio_embed, idx):
+        clean_embed = None
+        for i in range(1, len(audio_embed)):
+            if i in idx:
+                continue
+            else:
+                clean_embed = (
+                    audio_embed[i].reshape(1, -1)
+                    if clean_embed is None
+                    else torch.concatenate([clean_embed, audio_embed[i].reshape(1, -1)])
+                )
+        torch.save(clean_embed, self.config.path_audio_embedding)
+
+
+def argparser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-c",
+        "--config",
+        type=str,
+        default="config.yaml",
+        help="File path for config file.",
+    )
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == "__main__":
+    args = argparser()
+    embeder = Embeder(args.config)
+    embeder.run()

src/scrape.py

@@ -0,0 +1,138 @@
+import argparse
+import glob
+import os
+import time
+import urllib
+
+import librosa
+import pandas as pd
+import requests
+import soundfile as sf
+from bs4 import BeautifulSoup
+from omegaconf import OmegaConf
+from pydub import AudioSegment
+from requests.exceptions import Timeout
+
+
+class Scraper:
+    def __init__(self, config):
+        self.base_url = "https://soundeffect-lab.info/"
+        self.df = pd.DataFrame([], columns=["filename", "title", "category", "url"])
+        self.idx = 0
+        self.config = OmegaConf.load(config)
+        self.setup()
+        os.makedirs(self.config.path_data, exist_ok=True)
+        self.history = []
+
+    def run(self):
+        self.all_get()
+        self.preprocess()
+
+    def setup(self):
+        try:
+            html = requests.get(self.base_url, timeout=5)
+        except Timeout:
+            raise ValueError("Time Out")
+        soup = BeautifulSoup(html.content, "html.parser")
+        tags = soup.select("a")
+        self.urls = []
+        self.categories = []
+        for tag in tags:
+            category = tag.text
+            url = tag.get("href")
+            if "/sound/" in url:
+                self.urls.append(url)
+                self.categories.append(category)
+
+    def all_get(self):
+        for i in range(len(self.urls)):
+            now_url = self.base_url + self.urls[i][1:]
+            self.download(now_url, self.categories[i])
+        self.df.to_csv(self.config.path_csv)
+
+    def download(self, now_url, category):
+        try:
+            html = requests.get(now_url, timeout=5)
+            soup = BeautifulSoup(html.content, "html.parser")
+            body = soup.find(id="wrap").find("main")
+            tags = body.find(id="playarea").select("a")
+            count = 0
+            for tag in tags:
+                name = tag.get("download")
+                url = tag.get("href")
+                filename = os.path.join(self.config.path_data, name)
+                if os.path.exists(filename):
+                    continue
+                try:
+                    urllib.request.urlretrieve(now_url + url, filename)
+                    title = name.replace(".mp3", "")
+                    self.df.loc[self.idx] = {
+                        "filename": filename,
+                        "title": title,
+                        "category": category,
+                        "url": f"https://soundeffect-lab.info/sound/search.php?s={title}",
+                    }
+                    self.idx += 1
+                    time.sleep(2)
+                    count += 1
+                except Exception:
+                    continue
+            self.history.append(category)
+            print(now_url, category, len(tags), count)
+            paths = glob.glob(os.path.join(self.config.path_data, "*"))
+            assert len(paths) == len(self.df)
+
+            others = body.find(id="pagemenu-top").select("a")
+            other_urls, other_categories = [], []
+            for other in others:
+                other_url = other.get("href")
+                other_name = other.find("img").get("alt")
+                if other_name in self.history:
+                    continue
+                other_urls.append(other_url)
+                other_categories.append(other_name)
+            for i in range(len(other_urls)):
+                self.download(self.base_url + other_urls[i][1:], other_categories[i])
+        except Timeout:
+            print(f"Time Out: {now_url}")
+
+    def preprocess(self):
+        for i in range(len(self.df)):
+            song = AudioSegment.from_mp3(
+                os.path.join(self.config.path_data, self.df.iloc[i]["filename"])
+            )
+            song.export(
+                os.path.join(
+                    self.config.path_data,
+                    self.df.iloc[i]["filename"].replace(".mp3", ".wav"),
+                ),
+                format="wav",
+            )
+
+        for i in range(len(self.df)):
+            file = os.path.join(
+                self.config.path_data,
+                self.df.iloc[i]["filename"].replace(".mp3", ".wav"),
+            )
+            y, sr = librosa.core.load(file, sr=self.config.sample_rate, mono=True)
+            dir, name = os.path.split(file)
+            sf.write(os.path.join(dir, "new_" + name), y, sr, subtype="PCM_16")
+
+
+def argparser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-c",
+        "--config",
+        type=str,
+        default="config.yaml",
+        help="File path for config file.",
+    )
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == "__main__":
+    args = argparser()
+    scraper = Scraper(args.config)
+    scraper.run()

src/search.py

@@ -0,0 +1,88 @@
+import librosa
+import numpy as np
+import pandas as pd
+import soundfile as sf
+import torch
+from omegaconf import OmegaConf
+from pydub import AudioSegment
+from transformers import (
+    AutoFeatureExtractor,
+    BertForSequenceClassification,
+    BertJapaneseTokenizer,
+    Wav2Vec2ForXVector,
+)
+
+
+class Search:
+    def __init__(self, config):
+        self.config = OmegaConf.load(config)
+        self.df = pd.read_csv(self.config.path_csv)[["title", "url"]]
+        self.audio_feature_extractor = AutoFeatureExtractor.from_pretrained(
+            "anton-l/wav2vec2-base-superb-sv"
+        )
+        self.audio_model = Wav2Vec2ForXVector.from_pretrained(
+            "anton-l/wav2vec2-base-superb-sv"
+        )
+        self.text_tokenizer = BertJapaneseTokenizer.from_pretrained(
+            "cl-tohoku/bert-base-japanese-whole-word-masking"
+        )
+        self.text_model = BertForSequenceClassification.from_pretrained(
+            "cl-tohoku/bert-base-japanese-whole-word-masking",
+            num_labels=2,
+            output_attentions=False,
+            output_hidden_states=True,
+        ).eval()
+        self.text_reference = torch.load(self.config.path_text_embedding)
+        self.audio_reference = torch.load(self.config.path_audio_embedding)
+        self.similarity = torch.nn.CosineSimilarity(dim=-1)
+
+    def search(self, text, audio, ratio, topk):
+        text_embed, audio_embed = self.get_embedding(text, audio)
+        if text_embed is not None and audio_embed is not None:
+            result = self.similarity(
+                text_embed, self.text_reference
+            ) * ratio + self.similarity(audio_embed, self.audio_reference) * (1 - ratio)
+        elif text_embed is not None:
+            result = self.similarity(text_embed, self.text_reference)
+        elif audio_embed is not None:
+            result = self.similarity(audio_embed, self.audio_reference)
+        else:
+            raise ValueError("Input text or upload audio file.")
+        rank = np.argsort(result.numpy())[::-1][0 : int(topk)]
+        return self.df.iloc[rank]
+
+    def get_embedding(self, text, audio):
+        text_embed = None if text == "" else self._get_text_embedding(text)
+        audio_embed = None if audio is None else self._get_audio_embedding(audio)
+        return text_embed, audio_embed
+
+    def _get_text_embedding(self, text):
+        tokenized_text = self.text_tokenizer.tokenize(text)
+        indexed_tokens = self.text_tokenizer.convert_tokens_to_ids(tokenized_text)
+        tokens_tensor = torch.tensor([indexed_tokens])
+        with torch.no_grad():
+            all_encoder_layers = self.text_model(tokens_tensor)
+        embedding = torch.mean(all_encoder_layers[1][-2][0], axis=0).reshape(1, -1)
+        return embedding
+
+    def _get_audio_embedding(self, audio):
+        audio = self.preprocess_audio(audio)
+        song = AudioSegment.from_wav(audio)
+        song = np.array(song.get_array_of_samples(), dtype="float")
+        inputs = self.audio_feature_extractor(
+            [song],
+            sampling_rate=self.config.sample_rate,
+            return_tensors="pt",
+            padding=True,
+        )
+        with torch.no_grad():
+            embedding = self.audio_model(**inputs).embeddings
+        return embedding
+
+    def preprocess_audio(self, audio):
+        sample_rate, data = audio
+        audio = "tmp.wav"
+        sf.write(file=audio, data=data, samplerate=sample_rate)
+        y, sr = librosa.core.load(audio, sr=self.config.sample_rate, mono=True)
+        sf.write(audio, y, sr, subtype="PCM_16")
+        return audio