未验证 提交 41d194d1 编写于 作者: H Hui Zhang 提交者: GitHub

Merge pull request #1547 from qingen/database-search

[vec] add demo for searching vectors base on MySQL and Milvus, t…
([简体中文](./README_cn.md)|English)
# Audio Searching
## Introduction
As the Internet continues to evolve, unstructured data such as emails, social media photos, live videos, and customer service voice calls have become increasingly common. If we want to process the data on a computer, we need to use embedding technology to transform the data into vector and store, index, and query it
However, when there is a large amount of data, such as hundreds of millions of audio tracks, it is more difficult to do a similarity search. The exhaustive method is feasible, but very time consuming. For this scenario, this demo will introduce how to build an audio similarity retrieval system using the open source vector database Milvus
Audio retrieval (speech, music, speaker, etc.) enables querying and finding similar sounds (or the same speaker) in a large amount of audio data. The audio similarity retrieval system can be used to identify similar sound effects, minimize intellectual property infringement, quickly retrieve the voice print library, and help enterprises control fraud and identity theft. Audio retrieval also plays an important role in the classification and statistical analysis of audio data
In this demo, you will learn how to build an audio retrieval system to retrieve similar sound snippets. The uploaded audio clips are converted into vector data using paddlespeech-based pre-training models (audio classification model, speaker recognition model, etc.) and stored in Milvus. Milvus automatically generates a unique ID for each vector, then stores the ID and the corresponding audio information (audio ID, audio speaker ID, etc.) in MySQL to complete the library construction. During retrieval, users upload test audio to obtain vector, and then conduct vector similarity search in Milvus. The retrieval result returned by Milvus is vector ID, and the corresponding audio information can be queried in MySQL by ID
![Workflow of an audio searching system](./img/audio_searching.png)
Note:this demo uses the [CN-Celeb](http://openslr.org/82/) dataset of at least 650,000 audio entries and 3000 speakers to build the audio vector library, which is then retrieved using a preset distance calculation. The dataset can also use other, Adjust as needed, e.g. Librispeech, VoxCeleb, UrbanSound, GloVe, MNIST, etc
## Usage
### 1. Prepare MySQL and Milvus services by docker-compose
The audio similarity search system requires Milvus, MySQL services. We can start these containers with one click through [docker-compose.yaml](./docker-compose.yaml), so please make sure you have [installed Docker Engine](https://docs.docker.com/engine/install/) and [Docker Compose](https://docs.docker.com/compose/install/) before running. then
```bash
docker-compose -f docker-compose.yaml up -d
```
Then you will see the that all containers are created:
```bash
Creating network "quick_deploy_app_net" with driver "bridge"
Creating milvus-minio ... done
Creating milvus-etcd ... done
Creating audio-mysql ... done
Creating milvus-standalone ... done
Creating audio-webclient ... done
```
And show all containers with `docker ps`, and you can use `docker logs audio-mysql` to get the logs of server container
```bash
CONTAINER ID IMAGE COMMAND CREATED STATUS PORTS NAMES
b2bcf279e599 milvusdb/milvus:v2.0.1 "/tini -- milvus run…" 22 hours ago Up 22 hours 0.0.0.0:19530->19530/tcp milvus-standalone
d8ef4c84e25c mysql:5.7 "docker-entrypoint.s…" 22 hours ago Up 22 hours 0.0.0.0:3306->3306/tcp, 33060/tcp audio-mysql
8fb501edb4f3 quay.io/coreos/etcd:v3.5.0 "etcd -advertise-cli…" 22 hours ago Up 22 hours 2379-2380/tcp milvus-etcd
ffce340b3790 minio/minio:RELEASE.2020-12-03T00-03-10Z "/usr/bin/docker-ent…" 22 hours ago Up 22 hours (healthy) 9000/tcp milvus-minio
15c84a506754 iregistry.baidu-int.com/paddlespeech/audio-search-client:1.0 "/bin/bash -c '/usr/…" 22 hours ago Up 22 hours (healthy) 0.0.0.0:8068->80/tcp audio-webclient
```
### 2. Start API Server
Then to start the system server, and it provides HTTP backend services.
- Install the Python packages
```bash
pip install -r requirements.txt
```
- Set configuration
```bash
vim src/config.py
```
Modify the parameters according to your own environment. Here listing some parameters that need to be set, for more information please refer to [config.py](./src/config.py).
| **Parameter** | **Description** | **Default setting** |
| ---------------- | ----------------------------------------------------- | ------------------- |
| MILVUS_HOST | The IP address of Milvus, you can get it by ifconfig. If running everything on one machine, most likely 127.0.0.1 | 127.0.0.1 |
| MILVUS_PORT | Port of Milvus. | 19530 |
| VECTOR_DIMENSION | Dimension of the vectors. | 2048 |
| MYSQL_HOST | The IP address of Mysql. | 127.0.0.1 |
| MYSQL_PORT | Port of Milvus. | 3306 |
| DEFAULT_TABLE | The milvus and mysql default collection name. | audio_table |
- Run the code
Then start the server with Fastapi.
```bash
export PYTHONPATH=$PYTHONPATH:./src
python src/main.py
```
Then you will see the Application is started:
```bash
INFO: Started server process [3949]
2022-03-07 17:39:14,864 | INFO | server.py | serve | 75 | Started server process [3949]
INFO: Waiting for application startup.
2022-03-07 17:39:14,865 | INFO | on.py | startup | 45 | Waiting for application startup.
INFO: Application startup complete.
2022-03-07 17:39:14,866 | INFO | on.py | startup | 59 | Application startup complete.
INFO: Uvicorn running on http://0.0.0.0:8002 (Press CTRL+C to quit)
2022-03-07 17:39:14,867 | INFO | server.py | _log_started_message | 206 | Uvicorn running on http://0.0.0.0:8002 (Press CTRL+C to quit)
```
### 3. Usage
- Prepare data
```bash
wget -c https://www.openslr.org/resources/82/cn-celeb_v2.tar.gz && tar -xvf cn-celeb_v2.tar.gz
```
Note: If you want to build a quick demo, you can use ./src/test_main.py:download_audio_data function, it downloads 20 audio files , Subsequent results show this collection as an example
- scripts test (recommend!)
The internal process is downloading data, loading the Paddlespeech model, extracting embedding, storing library, retrieving and deleting library
```bash
python ./src/test_main.py
```
Output:
```bash
Checkpoint path: %your model path%
Extracting feature from audio No. 1 , 20 audios in total
Extracting feature from audio No. 2 , 20 audios in total
...
2022-03-09 17:22:13,870 | INFO | main.py | load_audios | 85 | Successfully loaded data, total count: 20
2022-03-09 17:22:13,898 | INFO | main.py | count_audio | 147 | Successfully count the number of data!
2022-03-09 17:22:13,918 | INFO | main.py | audio_path | 57 | Successfully load audio: ./example_audio/test.wav
...
2022-03-09 17:22:32,580 | INFO | main.py | search_local_audio | 131 | search result http://testserver/data?audio_path=./example_audio/test.wav, distance 0.0
2022-03-09 17:22:32,580 | INFO | main.py | search_local_audio | 131 | search result http://testserver/data?audio_path=./example_audio/knife_chopping.wav, distance 0.021805256605148315
2022-03-09 17:22:32,580 | INFO | main.py | search_local_audio | 131 | search result http://testserver/data?audio_path=./example_audio/knife_cut_into_flesh.wav, distance 0.052762262523174286
...
2022-03-09 17:22:32,582 | INFO | main.py | search_local_audio | 135 | Successfully searched similar audio!
2022-03-09 17:22:33,658 | INFO | main.py | drop_tables | 159 | Successfully drop tables in Milvus and MySQL!
```
- GUI test (optional)
Navigate to 127.0.0.1:8068 in your browser to access the front-end interface
Note: If the browser and the service are not on the same machine, then the IP needs to be changed to the IP of the machine where the service is located, and the corresponding API_URL in docker-compose.yaml needs to be changed and the service can be restarted
- Insert data
Download the data and decompress it to a path named /home/speech/data. Then enter /home/speech/data in the address bar of the upload page to upload the data
![](./img/insert.png)
- Search for similar audio
Select the magnifying glass icon on the left side of the interface. Then, press the "Default Target Audio File" button and upload a .wav sound file you'd like to search. Results will be displayed
![](./img/search.png)
### 4.Result
machine configuration:
- OS: CentOS release 7.6
- kernel:4.17.11-1.el7.elrepo.x86_64
- CPU:Intel(R) Xeon(R) CPU E5-2620 v4 @ 2.10GHz
- memory:132G
dataset:
- CN-Celeb, train size 650,000, test size 10,000, dimention 192, distance L2
recall and elapsed time statistics are shown in the following figure:
![](./img/result.png)
The retrieval framework based on Milvus takes about 2.9 milliseconds to retrieve on the premise of 90% recall rate, and it takes about 500 milliseconds for feature extraction (testing audio takes about 5 seconds), that is, a single audio test takes about 503 milliseconds in total, which can meet most application scenarios
### 5.Pretrained Models
Here is a list of pretrained models released by PaddleSpeech :
| Model | Sample Rate
| :--- | :---:
| ecapa_tdnn | 16000
| panns_cnn6| 32000
| panns_cnn10| 32000
| panns_cnn14| 32000
(简体中文|[English](./README.md))
# 音频相似性检索
## 介绍
随着互联网不断发展,电子邮件、社交媒体照片、直播视频、客服语音等非结构化数据已经变得越来越普遍。如果想要使用计算机来处理这些数据,需要使用 embedding 技术将这些数据转化为向量 vector,然后进行存储、建索引、并查询
但是,当数据量很大,比如上亿条音频要做相似度搜索,就比较困难了。穷举法固然可行,但非常耗时。针对这种场景,该demo 将介绍如何使用开源向量数据库 Milvus 搭建音频相似度检索系统
音频检索(如演讲、音乐、说话人等检索)实现了在海量音频数据中查询并找出相似声音(或相同说话人)片段。音频相似性检索系统可用于识别相似的音效、最大限度减少知识产权侵权等,还可以快速的检索声纹库、帮助企业控制欺诈和身份盗用等。在音频数据的分类和统计分析中,音频检索也发挥着重要作用
在本 demo 中,你将学会如何构建一个音频检索系统,用来检索相似的声音片段。使用基于 PaddleSpeech 预训练模型(音频分类模型,说话人识别模型等)将上传的音频片段转换为向量数据,并存储在 Milvus 中。Milvus 自动为每个向量生成唯一的 ID,然后将 ID 和 相应的音频信息(音频id,音频的说话人id等等)存储在 MySQL,这样就完成建库的工作。用户在检索时,上传测试音频,得到向量,然后在 Milvus 中进行向量相似度搜索,Milvus 返回的检索结果为向量 ID,通过 ID 在 MySQL 内部查询相应的音频信息即可
![音频检索流程图](./img/audio_searching.png)
注:该 demo 使用 [CN-Celeb](http://openslr.org/82/) 数据集,包括至少 650000 条音频,3000 个说话人,来建立音频向量库(音频特征,或音频说话人特征),然后通过预设的距离计算方式进行音频(或说话人)检索,这里面数据集也可以使用其他的,根据需要调整,如Librispeech,VoxCeleb,UrbanSound,GloVe,MNIST等
## 使用方法
### 1. MySQL 和 Milvus 安装
音频相似度搜索系统需要用到 Milvus, MySQL 服务。 我们可以通过 [docker-compose.yaml](./docker-compose.yaml) 一键启动这些容器,所以请确保在运行之前已经安装了 [Docker Engine](https://docs.docker.com/engine/install/)[Docker Compose](https://docs.docker.com/compose/install/)。 即
```bash
docker-compose -f docker-compose.yaml up -d
```
然后你会看到所有的容器都被创建:
```bash
Creating network "quick_deploy_app_net" with driver "bridge"
Creating milvus-minio ... done
Creating milvus-etcd ... done
Creating audio-mysql ... done
Creating milvus-standalone ... done
Creating audio-webclient ... done
```
可以采用'docker ps'来显示所有的容器,还可以使用'docker logs audio-mysql'来获取服务器容器的日志:
```bash
CONTAINER ID IMAGE COMMAND CREATED STATUS PORTS NAMES
b2bcf279e599 milvusdb/milvus:v2.0.1 "/tini -- milvus run…" 22 hours ago Up 22 hours 0.0.0.0:19530->19530/tcp milvus-standalone
d8ef4c84e25c mysql:5.7 "docker-entrypoint.s…" 22 hours ago Up 22 hours 0.0.0.0:3306->3306/tcp, 33060/tcp audio-mysql
8fb501edb4f3 quay.io/coreos/etcd:v3.5.0 "etcd -advertise-cli…" 22 hours ago Up 22 hours 2379-2380/tcp milvus-etcd
ffce340b3790 minio/minio:RELEASE.2020-12-03T00-03-10Z "/usr/bin/docker-ent…" 22 hours ago Up 22 hours (healthy) 9000/tcp milvus-minio
15c84a506754 iregistry.baidu-int.com/paddlespeech/audio-search-client:1.0 "/bin/bash -c '/usr/…" 22 hours ago Up 22 hours (healthy) 0.0.0.0:8068->80/tcp audio-webclient
```
### 2. 配置并启动 API 服务
启动系统服务程序,它会提供基于 Http 后端服务
- 安装服务依赖的 python 基础包
```bash
pip install -r requirements.txt
```
- 修改配置
```bash
vim src/config.py
```
请根据实际环境进行修改。 这里列出了一些需要设置的参数,更多信息请参考 [config.py](./src/config.py)
| **Parameter** | **Description** | **Default setting** |
| ---------------- | ----------------------------------------------------- | ------------------- |
| MILVUS_HOST | The IP address of Milvus, you can get it by ifconfig. If running everything on one machine, most likely 127.0.0.1 | 127.0.0.1 |
| MILVUS_PORT | Port of Milvus. | 19530 |
| VECTOR_DIMENSION | Dimension of the vectors. | 2048 |
| MYSQL_HOST | The IP address of Mysql. | 127.0.0.1 |
| MYSQL_PORT | Port of Milvus. | 3306 |
| DEFAULT_TABLE | The milvus and mysql default collection name. | audio_table |
- 运行程序
启动用 Fastapi 构建的服务
```bash
export PYTHONPATH=$PYTHONPATH:./src
python src/main.py
```
然后你会看到应用程序启动:
```bash
INFO: Started server process [3949]
2022-03-07 17:39:14,864 | INFO | server.py | serve | 75 | Started server process [3949]
INFO: Waiting for application startup.
2022-03-07 17:39:14,865 | INFO | on.py | startup | 45 | Waiting for application startup.
INFO: Application startup complete.
2022-03-07 17:39:14,866 | INFO | on.py | startup | 59 | Application startup complete.
INFO: Uvicorn running on http://0.0.0.0:8002 (Press CTRL+C to quit)
2022-03-07 17:39:14,867 | INFO | server.py | _log_started_message | 206 | Uvicorn running on http://0.0.0.0:8002 (Press CTRL+C to quit)
```
### 3. 测试方法
- 准备数据
```bash
wget -c https://www.openslr.org/resources/82/cn-celeb_v2.tar.gz && tar -xvf cn-celeb_v2.tar.gz
```
注:如果希望快速搭建 demo,可以采用 ./src/test_main.py:download_audio_data 内部的 20 条音频,另外后续结果展示以该集合为例
- 脚本测试(推荐)
```bash
python ./src/test_main.py
```
注:内部将依次下载数据,加载 paddlespeech 模型,提取 embedding,存储建库,检索,删库
输出:
```bash
Checkpoint path: %your model path%
Extracting feature from audio No. 1 , 20 audios in total
Extracting feature from audio No. 2 , 20 audios in total
...
2022-03-09 17:22:13,870 | INFO | main.py | load_audios | 85 | Successfully loaded data, total count: 20
2022-03-09 17:22:13,898 | INFO | main.py | count_audio | 147 | Successfully count the number of data!
2022-03-09 17:22:13,918 | INFO | main.py | audio_path | 57 | Successfully load audio: ./example_audio/test.wav
...
2022-03-09 17:22:32,580 | INFO | main.py | search_local_audio | 131 | search result http://testserver/data?audio_path=./example_audio/test.wav, distance 0.0
2022-03-09 17:22:32,580 | INFO | main.py | search_local_audio | 131 | search result http://testserver/data?audio_path=./example_audio/knife_chopping.wav, distance 0.021805256605148315
2022-03-09 17:22:32,580 | INFO | main.py | search_local_audio | 131 | search result http://testserver/data?audio_path=./example_audio/knife_cut_into_flesh.wav, distance 0.052762262523174286
...
2022-03-09 17:22:32,582 | INFO | main.py | search_local_audio | 135 | Successfully searched similar audio!
2022-03-09 17:22:33,658 | INFO | main.py | drop_tables | 159 | Successfully drop tables in Milvus and MySQL!
```
- 前端测试(可选)
在浏览器中输入 127.0.0.1:8068 访问前端页面
注:如果浏览器和服务不在同一台机器上,那么 IP 需要修改成服务所在的机器 IP,并且docker-compose.yaml 中相应的 API_URL 也要修改,并重新起服务即可
- 上传音频
下载数据并解压到一文件夹,假设为 /home/speech/data,那么在上传页面地址栏输入 /home/speech/data 进行数据上传
![](./img/insert.png)
- 检索相似音频
选择左上角放大镜,点击 “Default Target Audio File” 按钮,上传测试音频,接着你将看到检索结果
![](./img/search.png)
### 4. 结果
机器配置:
- 操作系统: CentOS release 7.6
- 内核:4.17.11-1.el7.elrepo.x86_64
- 处理器:Intel(R) Xeon(R) CPU E5-2620 v4 @ 2.10GHz
- 内存:132G
数据集:
- CN-Celeb, 训练集 65万, 测试集 1万,向量维度 192,距离计算方式 L2
召回和耗时统计如下图:
![](./img/result.png)
基于 milvus 的检索框架在召回率 90% 的前提下,检索耗时约 2.9 毫秒,加上特征提取(Embedding)耗时约 500毫秒(测试音频时长约 5秒),即单条音频测试总共耗时约 503 毫秒,可以满足大多数应用场景
### 5. 预训练模型
以下是 PaddleSpeech 提供的预训练模型列表:
| 模型 | 采样率
| :--- | :---:
| ecapa_tdnn| 16000
| panns_cnn6| 32000
| panns_cnn10| 32000
| panns_cnn14| 32000
version: '3.5'
services:
etcd:
container_name: milvus-etcd
image: quay.io/coreos/etcd:v3.5.0
networks:
app_net:
environment:
- ETCD_AUTO_COMPACTION_MODE=revision
- ETCD_AUTO_COMPACTION_RETENTION=1000
- ETCD_QUOTA_BACKEND_BYTES=4294967296
volumes:
- ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/etcd:/etcd
command: etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd
minio:
container_name: milvus-minio
image: minio/minio:RELEASE.2020-12-03T00-03-10Z
networks:
app_net:
environment:
MINIO_ACCESS_KEY: minioadmin
MINIO_SECRET_KEY: minioadmin
volumes:
- ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/minio:/minio_data
command: minio server /minio_data
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:9000/minio/health/live"]
interval: 30s
timeout: 20s
retries: 3
standalone:
container_name: milvus-standalone
image: milvusdb/milvus:v2.0.1
networks:
app_net:
ipv4_address: 172.16.23.10
command: ["milvus", "run", "standalone"]
environment:
ETCD_ENDPOINTS: etcd:2379
MINIO_ADDRESS: minio:9000
volumes:
- ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/milvus:/var/lib/milvus
ports:
- "19530:19530"
depends_on:
- "etcd"
- "minio"
mysql:
container_name: audio-mysql
image: mysql:5.7
networks:
app_net:
ipv4_address: 172.16.23.11
environment:
- MYSQL_ROOT_PASSWORD=123456
volumes:
- ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/mysql:/var/lib/mysql
ports:
- "3306:3306"
webclient:
container_name: audio-webclient
image: qingen1/paddlespeech-audio-search-client:2.3
networks:
app_net:
ipv4_address: 172.16.23.13
environment:
API_URL: 'http://127.0.0.1:8002'
ports:
- "8068:80"
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost/"]
interval: 30s
timeout: 20s
retries: 3
networks:
app_net:
driver: bridge
ipam:
driver: default
config:
- subnet: 172.16.23.0/24
gateway: 172.16.23.1
soundfile==0.10.3.post1
librosa==0.8.0
numpy
pymysql
fastapi
uvicorn
diskcache==5.2.1
pymilvus==2.0.1
python-multipart
typing
starlette
pydantic
\ No newline at end of file
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed 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
#
# http://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
############### Milvus Configuration ###############
MILVUS_HOST = os.getenv("MILVUS_HOST", "127.0.0.1")
MILVUS_PORT = int(os.getenv("MILVUS_PORT", "19530"))
VECTOR_DIMENSION = int(os.getenv("VECTOR_DIMENSION", "2048"))
INDEX_FILE_SIZE = int(os.getenv("INDEX_FILE_SIZE", "1024"))
METRIC_TYPE = os.getenv("METRIC_TYPE", "L2")
DEFAULT_TABLE = os.getenv("DEFAULT_TABLE", "audio_table")
TOP_K = int(os.getenv("TOP_K", "10"))
############### MySQL Configuration ###############
MYSQL_HOST = os.getenv("MYSQL_HOST", "127.0.0.1")
MYSQL_PORT = int(os.getenv("MYSQL_PORT", "3306"))
MYSQL_USER = os.getenv("MYSQL_USER", "root")
MYSQL_PWD = os.getenv("MYSQL_PWD", "123456")
MYSQL_DB = os.getenv("MYSQL_DB", "mysql")
############### Data Path ###############
UPLOAD_PATH = os.getenv("UPLOAD_PATH", "tmp/audio-data")
############### Number of Log Files ###############
LOGS_NUM = int(os.getenv("logs_num", "0"))
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed 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
#
# http://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 librosa
import numpy as np
from logs import LOGGER
def get_audio_embedding(path):
"""
Use vpr_inference to generate embedding of audio
"""
try:
RESAMPLE_RATE = 16000
audio, _ = librosa.load(path, sr=RESAMPLE_RATE, mono=True)
# TODO add infer/python interface to get embedding, now fake it by rand
# vpr = ECAPATDNN(checkpoint_path=None, device='cuda')
# embedding = vpr.inference(audio)
np.random.seed(hash(os.path.basename(path)) % 1000000)
embedding = np.random.rand(1, 2048)
embedding = embedding / np.linalg.norm(embedding)
embedding = embedding.tolist()[0]
return embedding
except Exception as e:
LOGGER.error(f"Error with embedding:{e}")
return None
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed 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
#
# http://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 codecs
import datetime
import logging
import os
import re
import sys
from config import LOGS_NUM
class MultiprocessHandler(logging.FileHandler):
"""
A handler class which writes formatted logging records to disk files
"""
def __init__(self,
filename,
when='D',
backupCount=0,
encoding=None,
delay=False):
"""
Open the specified file and use it as the stream for logging
"""
self.prefix = filename
self.backupCount = backupCount
self.when = when.upper()
self.extMath = r"^\d{4}-\d{2}-\d{2}"
self.when_dict = {
'S': "%Y-%m-%d-%H-%M-%S",
'M': "%Y-%m-%d-%H-%M",
'H': "%Y-%m-%d-%H",
'D': "%Y-%m-%d"
}
self.suffix = self.when_dict.get(when)
if not self.suffix:
print('The specified date interval unit is invalid: ', self.when)
sys.exit(1)
self.filefmt = os.path.join('.', "logs",
f"{self.prefix}-{self.suffix}.log")
self.filePath = datetime.datetime.now().strftime(self.filefmt)
_dir = os.path.dirname(self.filefmt)
try:
if not os.path.exists(_dir):
os.makedirs(_dir)
except Exception as e:
print('Failed to create log file: ', e)
print("log_path:" + self.filePath)
sys.exit(1)
logging.FileHandler.__init__(self, self.filePath, 'a+', encoding, delay)
def should_change_file_to_write(self):
"""
To write the file
"""
_filePath = datetime.datetime.now().strftime(self.filefmt)
if _filePath != self.filePath:
self.filePath = _filePath
return True
return False
def do_change_file(self):
"""
To change file states
"""
self.baseFilename = os.path.abspath(self.filePath)
if self.stream:
self.stream.close()
self.stream = None
if not self.delay:
self.stream = self._open()
if self.backupCount > 0:
for s in self.get_files_to_delete():
os.remove(s)
def get_files_to_delete(self):
"""
To delete backup files
"""
dir_name, _ = os.path.split(self.baseFilename)
file_names = os.listdir(dir_name)
result = []
prefix = self.prefix + '-'
for file_name in file_names:
if file_name[:len(prefix)] == prefix:
suffix = file_name[len(prefix):-4]
if re.compile(self.extMath).match(suffix):
result.append(os.path.join(dir_name, file_name))
result.sort()
if len(result) < self.backupCount:
result = []
else:
result = result[:len(result) - self.backupCount]
return result
def emit(self, record):
"""
Emit a record
"""
try:
if self.should_change_file_to_write():
self.do_change_file()
logging.FileHandler.emit(self, record)
except (KeyboardInterrupt, SystemExit):
raise
except:
self.handleError(record)
def write_log():
"""
Init a logger
"""
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
# formatter = '%(asctime)s | %(levelname)s | %(filename)s | %(funcName)s | %(module)s | %(lineno)s | %(message)s'
fmt = logging.Formatter(
'%(asctime)s | %(levelname)s | %(filename)s | %(funcName)s | %(lineno)s | %(message)s'
)
stream_handler = logging.StreamHandler(sys.stdout)
stream_handler.setLevel(logging.INFO)
stream_handler.setFormatter(fmt)
log_name = "audio-searching"
file_handler = MultiprocessHandler(log_name, when='D', backupCount=LOGS_NUM)
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(fmt)
file_handler.do_change_file()
logger.addHandler(stream_handler)
logger.addHandler(file_handler)
return logger
LOGGER = write_log()
if __name__ == "__main__":
message = 'test writing logs'
LOGGER.info(message)
LOGGER.debug(message)
LOGGER.error(message)
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed 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
#
# http://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
from typing import Optional
import uvicorn
from config import UPLOAD_PATH
from diskcache import Cache
from fastapi import FastAPI
from fastapi import File
from fastapi import UploadFile
from logs import LOGGER
from milvus_helpers import MilvusHelper
from mysql_helpers import MySQLHelper
from operations.count import do_count
from operations.drop import do_drop
from operations.load import do_load
from operations.search import do_search
from pydantic import BaseModel
from starlette.middleware.cors import CORSMiddleware
from starlette.requests import Request
from starlette.responses import FileResponse
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"])
MODEL = None
MILVUS_CLI = MilvusHelper()
MYSQL_CLI = MySQLHelper()
# Mkdir 'tmp/audio-data'
if not os.path.exists(UPLOAD_PATH):
os.makedirs(UPLOAD_PATH)
LOGGER.info(f"Mkdir the path: {UPLOAD_PATH}")
@app.get('/data')
def audio_path(audio_path):
# Get the audio file
try:
LOGGER.info(f"Successfully load audio: {audio_path}")
return FileResponse(audio_path)
except Exception as e:
LOGGER.error(f"upload audio error: {e}")
return {'status': False, 'msg': e}, 400
@app.get('/progress')
def get_progress():
# Get the progress of dealing with data
try:
cache = Cache('./tmp')
return f"current: {cache['current']}, total: {cache['total']}"
except Exception as e:
LOGGER.error(f"Upload data error: {e}")
return {'status': False, 'msg': e}, 400
class Item(BaseModel):
Table: Optional[str] = None
File: str
@app.post('/audio/load')
async def load_audios(item: Item):
# Insert all the audio files under the file path to Milvus/MySQL
try:
total_num = do_load(item.Table, item.File, MILVUS_CLI, MYSQL_CLI)
LOGGER.info(f"Successfully loaded data, total count: {total_num}")
return {'status': True, 'msg': "Successfully loaded data!"}
except Exception as e:
LOGGER.error(e)
return {'status': False, 'msg': e}, 400
@app.post('/audio/search')
async def search_audio(request: Request,
table_name: str=None,
audio: UploadFile=File(...)):
# Search the uploaded audio in Milvus/MySQL
try:
# Save the upload data to server.
content = await audio.read()
query_audio_path = os.path.join(UPLOAD_PATH, audio.filename)
with open(query_audio_path, "wb+") as f:
f.write(content)
host = request.headers['host']
_, paths, distances = do_search(host, table_name, query_audio_path,
MILVUS_CLI, MYSQL_CLI)
names = []
for path, score in zip(paths, distances):
names.append(os.path.basename(path))
LOGGER.info(f"search result {path}, score {score}")
res = dict(zip(paths, zip(names, distances)))
# Sort results by distance metric, closest distances first
res = sorted(res.items(), key=lambda item: item[1][1], reverse=True)
LOGGER.info("Successfully searched similar audio!")
return res
except Exception as e:
LOGGER.error(e)
return {'status': False, 'msg': e}, 400
@app.post('/audio/search/local')
async def search_local_audio(request: Request,
query_audio_path: str,
table_name: str=None):
# Search the uploaded audio in Milvus/MySQL
try:
host = request.headers['host']
_, paths, distances = do_search(host, table_name, query_audio_path,
MILVUS_CLI, MYSQL_CLI)
names = []
for path, score in zip(paths, distances):
names.append(os.path.basename(path))
LOGGER.info(f"search result {path}, score {score}")
res = dict(zip(paths, zip(names, distances)))
# Sort results by distance metric, closest distances first
res = sorted(res.items(), key=lambda item: item[1][1], reverse=True)
LOGGER.info("Successfully searched similar audio!")
return res
except Exception as e:
LOGGER.error(e)
return {'status': False, 'msg': e}, 400
@app.get('/audio/count')
async def count_audio(table_name: str=None):
# Returns the total number of vectors in the system
try:
num = do_count(table_name, MILVUS_CLI)
LOGGER.info("Successfully count the number of data!")
return num
except Exception as e:
LOGGER.error(e)
return {'status': False, 'msg': e}, 400
@app.post('/audio/drop')
async def drop_tables(table_name: str=None):
# Delete the collection of Milvus and MySQL
try:
status = do_drop(table_name, MILVUS_CLI, MYSQL_CLI)
LOGGER.info("Successfully drop tables in Milvus and MySQL!")
return status
except Exception as e:
LOGGER.error(e)
return {'status': False, 'msg': e}, 400
if __name__ == '__main__':
uvicorn.run(app=app, host='0.0.0.0', port=8002)
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed 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
#
# http://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 sys
from config import METRIC_TYPE
from config import MILVUS_HOST
from config import MILVUS_PORT
from config import VECTOR_DIMENSION
from logs import LOGGER
from pymilvus import Collection
from pymilvus import CollectionSchema
from pymilvus import connections
from pymilvus import DataType
from pymilvus import FieldSchema
from pymilvus import utility
class MilvusHelper:
"""
the basic operations of PyMilvus
# This example shows how to:
# 1. connect to Milvus server
# 2. create a collection
# 3. insert entities
# 4. create index
# 5. search
# 6. delete a collection
"""
def __init__(self):
try:
self.collection = None
connections.connect(host=MILVUS_HOST, port=MILVUS_PORT)
LOGGER.debug(
f"Successfully connect to Milvus with IP:{MILVUS_HOST} and PORT:{MILVUS_PORT}"
)
except Exception as e:
LOGGER.error(f"Failed to connect Milvus: {e}")
sys.exit(1)
def set_collection(self, collection_name):
try:
if self.has_collection(collection_name):
self.collection = Collection(name=collection_name)
else:
raise Exception(
f"There is no collection named:{collection_name}")
except Exception as e:
LOGGER.error(f"Failed to set collection in Milvus: {e}")
sys.exit(1)
def has_collection(self, collection_name):
# Return if Milvus has the collection
try:
return utility.has_collection(collection_name)
except Exception as e:
LOGGER.error(f"Failed to check state of collection in Milvus: {e}")
sys.exit(1)
def create_collection(self, collection_name):
# Create milvus collection if not exists
try:
if not self.has_collection(collection_name):
field1 = FieldSchema(
name="id",
dtype=DataType.INT64,
descrition="int64",
is_primary=True,
auto_id=True)
field2 = FieldSchema(
name="embedding",
dtype=DataType.FLOAT_VECTOR,
descrition="speaker embeddings",
dim=VECTOR_DIMENSION,
is_primary=False)
schema = CollectionSchema(
fields=[field1, field2], description="embeddings info")
self.collection = Collection(
name=collection_name, schema=schema)
LOGGER.debug(f"Create Milvus collection: {collection_name}")
else:
self.set_collection(collection_name)
return "OK"
except Exception as e:
LOGGER.error(f"Failed to create collection in Milvus: {e}")
sys.exit(1)
def insert(self, collection_name, vectors):
# Batch insert vectors to milvus collection
try:
self.create_collection(collection_name)
data = [vectors]
self.set_collection(collection_name)
mr = self.collection.insert(data)
ids = mr.primary_keys
self.collection.load()
LOGGER.debug(
f"Insert vectors to Milvus in collection: {collection_name} with {len(vectors)} rows"
)
return ids
except Exception as e:
LOGGER.error(f"Failed to insert data to Milvus: {e}")
sys.exit(1)
def create_index(self, collection_name):
# Create IVF_FLAT index on milvus collection
try:
self.set_collection(collection_name)
default_index = {
"index_type": "IVF_SQ8",
"metric_type": METRIC_TYPE,
"params": {
"nlist": 16384
}
}
status = self.collection.create_index(
field_name="embedding", index_params=default_index)
if not status.code:
LOGGER.debug(
f"Successfully create index in collection:{collection_name} with param:{default_index}"
)
return status
else:
raise Exception(status.message)
except Exception as e:
LOGGER.error(f"Failed to create index: {e}")
sys.exit(1)
def delete_collection(self, collection_name):
# Delete Milvus collection
try:
self.set_collection(collection_name)
self.collection.drop()
LOGGER.debug("Successfully drop collection!")
return "ok"
except Exception as e:
LOGGER.error(f"Failed to drop collection: {e}")
sys.exit(1)
def search_vectors(self, collection_name, vectors, top_k):
# Search vector in milvus collection
try:
self.set_collection(collection_name)
search_params = {
"metric_type": METRIC_TYPE,
"params": {
"nprobe": 16
}
}
res = self.collection.search(
vectors,
anns_field="embedding",
param=search_params,
limit=top_k)
LOGGER.debug(f"Successfully search in collection: {res}")
return res
except Exception as e:
LOGGER.error(f"Failed to search vectors in Milvus: {e}")
sys.exit(1)
def count(self, collection_name):
# Get the number of milvus collection
try:
self.set_collection(collection_name)
num = self.collection.num_entities
LOGGER.debug(
f"Successfully get the num:{num} of the collection:{collection_name}"
)
return num
except Exception as e:
LOGGER.error(f"Failed to count vectors in Milvus: {e}")
sys.exit(1)
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed 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
#
# http://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 sys
import pymysql
from config import MYSQL_DB
from config import MYSQL_HOST
from config import MYSQL_PORT
from config import MYSQL_PWD
from config import MYSQL_USER
from logs import LOGGER
class MySQLHelper():
"""
the basic operations of PyMySQL
# This example shows how to:
# 1. connect to MySQL server
# 2. create a table
# 3. insert data to table
# 4. search by milvus ids
# 5. delete table
"""
def __init__(self):
self.conn = pymysql.connect(
host=MYSQL_HOST,
user=MYSQL_USER,
port=MYSQL_PORT,
password=MYSQL_PWD,
database=MYSQL_DB,
local_infile=True)
self.cursor = self.conn.cursor()
def test_connection(self):
try:
self.conn.ping()
except Exception:
self.conn = pymysql.connect(
host=MYSQL_HOST,
user=MYSQL_USER,
port=MYSQL_PORT,
password=MYSQL_PWD,
database=MYSQL_DB,
local_infile=True)
self.cursor = self.conn.cursor()
def create_mysql_table(self, table_name):
# Create mysql table if not exists
self.test_connection()
sql = "create table if not exists " + table_name + "(milvus_id TEXT, audio_path TEXT);"
try:
self.cursor.execute(sql)
LOGGER.debug(f"MYSQL create table: {table_name} with sql: {sql}")
except Exception as e:
LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
sys.exit(1)
def load_data_to_mysql(self, table_name, data):
# Batch insert (Milvus_ids, img_path) to mysql
self.test_connection()
sql = "insert into " + table_name + " (milvus_id,audio_path) values (%s,%s);"
try:
self.cursor.executemany(sql, data)
self.conn.commit()
LOGGER.debug(
f"MYSQL loads data to table: {table_name} successfully")
except Exception as e:
LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
sys.exit(1)
def search_by_milvus_ids(self, ids, table_name):
# Get the img_path according to the milvus ids
self.test_connection()
str_ids = str(ids).replace('[', '').replace(']', '')
sql = "select audio_path from " + table_name + " where milvus_id in (" + str_ids + ") order by field (milvus_id," + str_ids + ");"
try:
self.cursor.execute(sql)
results = self.cursor.fetchall()
results = [res[0] for res in results]
LOGGER.debug("MYSQL search by milvus id.")
return results
except Exception as e:
LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
sys.exit(1)
def delete_table(self, table_name):
# Delete mysql table if exists
self.test_connection()
sql = "drop table if exists " + table_name + ";"
try:
self.cursor.execute(sql)
LOGGER.debug(f"MYSQL delete table:{table_name}")
except Exception as e:
LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
sys.exit(1)
def delete_all_data(self, table_name):
# Delete all the data in mysql table
self.test_connection()
sql = 'delete from ' + table_name + ';'
try:
self.cursor.execute(sql)
self.conn.commit()
LOGGER.debug(f"MYSQL delete all data in table:{table_name}")
except Exception as e:
LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
sys.exit(1)
def count_table(self, table_name):
# Get the number of mysql table
self.test_connection()
sql = "select count(milvus_id) from " + table_name + ";"
try:
self.cursor.execute(sql)
results = self.cursor.fetchall()
LOGGER.debug(f"MYSQL count table:{table_name}")
return results[0][0]
except Exception as e:
LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
sys.exit(1)
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed 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
#
# http://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.
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed 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
#
# http://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 sys
from config import DEFAULT_TABLE
from logs import LOGGER
def do_count(table_name, milvus_cli):
"""
Returns the total number of vectors in the system
"""
if not table_name:
table_name = DEFAULT_TABLE
try:
if not milvus_cli.has_collection(table_name):
return None
num = milvus_cli.count(table_name)
return num
except Exception as e:
LOGGER.error(f"Error attempting to count table {e}")
sys.exit(1)
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed 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
#
# http://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 sys
from config import DEFAULT_TABLE
from logs import LOGGER
def do_drop(table_name, milvus_cli, mysql_cli):
"""
Delete the collection of Milvus and MySQL
"""
if not table_name:
table_name = DEFAULT_TABLE
try:
if not milvus_cli.has_collection(table_name):
return "Collection is not exist"
status = milvus_cli.delete_collection(table_name)
mysql_cli.delete_table(table_name)
return status
except Exception as e:
LOGGER.error(f"Error attempting to drop table: {e}")
sys.exit(1)
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed 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
#
# http://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 sys
from config import DEFAULT_TABLE
from diskcache import Cache
from encode import get_audio_embedding
from logs import LOGGER
def get_audios(path):
"""
List all wav and aif files recursively under the path folder.
"""
supported_formats = [".wav", ".mp3", ".ogg", ".flac", ".m4a"]
return [
item
for sublist in [[os.path.join(dir, file) for file in files]
for dir, _, files in list(os.walk(path))]
for item in sublist if os.path.splitext(item)[1] in supported_formats
]
def extract_features(audio_dir):
"""
Get the vector of audio
"""
try:
cache = Cache('./tmp')
feats = []
names = []
audio_list = get_audios(audio_dir)
total = len(audio_list)
cache['total'] = total
for i, audio_path in enumerate(audio_list):
norm_feat = get_audio_embedding(audio_path)
if norm_feat is None:
continue
feats.append(norm_feat)
names.append(audio_path.encode())
cache['current'] = i + 1
print(
f"Extracting feature from audio No. {i + 1} , {total} audios in total"
)
return feats, names
except Exception as e:
LOGGER.error(f"Error with extracting feature from audio {e}")
sys.exit(1)
def format_data(ids, names):
"""
Combine the id of the vector and the name of the audio into a list
"""
data = []
for i in range(len(ids)):
value = (str(ids[i]), names[i])
data.append(value)
return data
def do_load(table_name, audio_dir, milvus_cli, mysql_cli):
"""
Import vectors to Milvus and data to Mysql respectively
"""
if not table_name:
table_name = DEFAULT_TABLE
vectors, names = extract_features(audio_dir)
ids = milvus_cli.insert(table_name, vectors)
milvus_cli.create_index(table_name)
mysql_cli.create_mysql_table(table_name)
mysql_cli.load_data_to_mysql(table_name, format_data(ids, names))
return len(ids)
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed 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
#
# http://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 sys
from config import DEFAULT_TABLE
from config import TOP_K
from encode import get_audio_embedding
from logs import LOGGER
def do_search(host, table_name, audio_path, milvus_cli, mysql_cli):
"""
Search the uploaded audio in Milvus/MySQL
"""
try:
if not table_name:
table_name = DEFAULT_TABLE
feat = get_audio_embedding(audio_path)
vectors = milvus_cli.search_vectors(table_name, [feat], TOP_K)
vids = [str(x.id) for x in vectors[0]]
paths = mysql_cli.search_by_milvus_ids(vids, table_name)
distances = [x.distance for x in vectors[0]]
for i in range(len(paths)):
tmp = "http://" + str(host) + "/data?audio_path=" + str(paths[i])
paths[i] = tmp
distances[i] = (1 - distances[i]) * 100
return vids, paths, distances
except Exception as e:
LOGGER.error(f"Error with search: {e}")
sys.exit(1)
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed 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
#
# http://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 zipfile
import gdown
from fastapi.testclient import TestClient
from main import app
client = TestClient(app)
def download_audio_data():
"""
download audio data
"""
url = 'https://drive.google.com/uc?id=1bKu21JWBfcZBuEuzFEvPoAX6PmRrgnUp'
gdown.download(url)
with zipfile.ZipFile('example_audio.zip', 'r') as zip_ref:
zip_ref.extractall('./example_audio')
def test_drop():
"""
Delete the collection of Milvus and MySQL
"""
response = client.post("/audio/drop")
assert response.status_code == 200
def test_load():
"""
Insert all the audio files under the file path to Milvus/MySQL
"""
response = client.post("/audio/load", json={"File": "./example_audio"})
assert response.status_code == 200
assert response.json() == {
'status': True,
'msg': "Successfully loaded data!"
}
def test_progress():
"""
Get the progress of dealing with data
"""
response = client.get("/progress")
assert response.status_code == 200
assert response.json() == "current: 20, total: 20"
def test_count():
"""
Returns the total number of vectors in the system
"""
response = client.get("audio/count")
assert response.status_code == 200
assert response.json() == 20
def test_search():
"""
Search the uploaded audio in Milvus/MySQL
"""
response = client.post(
"/audio/search/local?query_audio_path=.%2Fexample_audio%2Ftest.wav")
assert response.status_code == 200
assert len(response.json()) == 10
def test_data():
"""
Get the audio file
"""
response = client.get("/data?audio_path=.%2Fexample_audio%2Ftest.wav")
assert response.status_code == 200
if __name__ == "__main__":
download_audio_data()
test_load()
test_count()
test_search()
test_drop()
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册