main.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
docker build -t template_match_asr:v20250109_1615 .

docker run -itd \
--name template_match_asr_7862 \
--restart=always \
--network host \
-e port=7862 \
template_match_asr:v20250109_1615
"""
import argparse
from pathlib import Path
import platform
import tempfile
import time
import uuid

import gradio as gr
import librosa
import numpy as np
from scipy.io import wavfile

from project_settings import environment, project_path

from tabs.annotation import get_annotation_tab


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--trained_model_dir",
        default=(project_path / "trained_models").as_posix(),
        type=str,
    )
    parser.add_argument("--win_size", default=2.0, type=int)
    parser.add_argument("--win_step", default=0.25, type=int)
    args = parser.parse_args()
    return args


def convert_sample_rate(signal: np.ndarray, sample_rate: int, target_sample_rate: int):
    filename = save_input_audio(sample_rate, signal)

    signal, _ = librosa.load(filename, sr=target_sample_rate)
    signal = np.array(signal * (1 << 15), dtype=np.int16)
    return signal


def save_input_audio(sample_rate: int, signal: np.ndarray) -> str:

    temp_audio_dir = Path(tempfile.gettempdir()) / "input_audio"
    temp_audio_dir.mkdir(parents=True, exist_ok=True)
    filename = temp_audio_dir / f"{uuid.uuid4()}.wav"
    filename = filename.as_posix()
    wavfile.write(
        filename,
        sample_rate, signal
    )

    return filename


def generate_short_id():
    timestamp = int(time.time() * 1000)  # 毫秒级时间戳
    return f"{timestamp}"[-5:]


def main():
    args = get_args()

    trained_model_dir = Path(args.trained_model_dir)

    # ui
    with gr.Blocks() as blocks:
        with gr.Tabs():
            _ = get_annotation_tab(
                trained_model_dir=trained_model_dir,
                win_size=args.win_size,
                win_step=args.win_step,
            )

    # http://127.0.0.1:7862/
    blocks.queue().launch(
        share=False if platform.system() == "Windows" else False,
        server_name="127.0.0.1" if platform.system() == "Windows" else "0.0.0.0",
        # server_name="0.0.0.0",
        server_port=environment.get("port", 7860, dtype=int),
        # allowed_paths=[
        #     "E:/Users/tianx/HuggingDatasets/international_voice/data/early_media"
        # ]
        show_error=True,
    )

    return


if __name__ == "__main__":
    main()