package main
|
|
import (
|
"bufio"
|
"bytes"
|
"encoding/base64"
|
"encoding/json"
|
"fmt"
|
"log"
|
"net/url"
|
"os"
|
"os/exec"
|
"strconv"
|
"strings"
|
"sync"
|
"time"
|
|
"github.com/gorilla/websocket"
|
)
|
|
var args struct {
|
audio_in string
|
thread_num int
|
host string
|
port string
|
output_dir string
|
hotword string
|
audio_fs int
|
use_itn int
|
mode string
|
chunk_size []int
|
chunk_interval int
|
}
|
|
var websocketConn *websocket.Conn
|
var offline_msg_done = false
|
|
type Message struct {
|
WavName string `json:"wav_name"`
|
Text string `json:"text"`
|
TimeStamp string `json:"timestamp"`
|
Mode string `json:"mode"`
|
}
|
|
type AudioData struct {
|
SampleRate int `json:"sample_rate"`
|
Stride int `json:"stride"`
|
ChunkNum int `json:"chunk_num"`
|
AudioBytes string `json:"audio_bytes"`
|
}
|
|
func IntSlicetoString(nums []int) string {
|
strNums := make([]string, len(nums))
|
for i, num := range nums {
|
strNums[i] = strconv.Itoa(num)
|
}
|
result := strings.Join(strNums, ",")
|
return result
|
}
|
|
func recordFromScp(chunk_begin, chunk_size int) {
|
wavs := []string{args.audio_in}
|
sample_rate := args.audio_fs
|
wav_format := "pcm"
|
use_itn := true
|
|
wav_name := "demo"
|
wav_path := wavs[0]
|
var audio_bytes []byte
|
var stride, chunk_num int
|
|
fst_dict := make(map[string]int)
|
hotword_msg := ""
|
if args.hotword != "" {
|
file, err := os.Open(args.hotword)
|
if err != nil {
|
log.Fatalf("failed to open file: %v", err)
|
}
|
defer file.Close()
|
|
scanner := bufio.NewScanner(file)
|
for scanner.Scan() {
|
words := strings.Fields(scanner.Text())
|
|
if len(words) < 2 {
|
fmt.Println("Please checkout format of hotwords")
|
continue
|
}
|
weight, err := strconv.Atoi(words[1])
|
if err != nil {
|
fmt.Println("The weight of hotwords must be Int!")
|
}
|
fst_dict[words[0]] = weight
|
}
|
if err := scanner.Err(); err != nil {
|
log.Fatalf("error reading file: %v", err)
|
}
|
bytes, _ := json.Marshal(fst_dict)
|
hotword_msg = string(bytes)
|
fmt.Println("HotWord: ", hotword_msg)
|
}
|
|
if args.use_itn == 0 {
|
use_itn = false
|
}
|
|
if chunk_size > 0 {
|
wavs = wavs[chunk_begin : chunk_begin+chunk_size]
|
}
|
|
if strings.HasSuffix(wav_path, ".wav") {
|
cmd := exec.Command("python", "wavhandler.py", wav_path, IntSlicetoString(args.chunk_size), strconv.Itoa(args.chunk_interval))
|
|
var out bytes.Buffer
|
cmd.Stdout = &out
|
err := cmd.Run()
|
if err != nil {
|
fmt.Println("Error running Python script:", err)
|
return
|
}
|
|
var audioData AudioData
|
err = json.Unmarshal(out.Bytes(), &audioData)
|
if err != nil {
|
fmt.Println("Error parsing JSON:", err)
|
return
|
}
|
|
stride = audioData.Stride
|
chunk_num = audioData.ChunkNum
|
sample_rate = audioData.SampleRate
|
audio_bytes, err = base64.StdEncoding.DecodeString(audioData.AudioBytes)
|
if err != nil {
|
fmt.Println("Error decoding Base64:", err)
|
return
|
}
|
} else {
|
fmt.Println("Currently, only the WAV format is supported")
|
return
|
}
|
|
first_message := make(map[string]interface{})
|
first_message["mode"] = args.mode
|
first_message["chunk_size"] = args.chunk_size
|
first_message["chunk_interval"] = args.chunk_interval
|
first_message["audio_fs"] = sample_rate
|
first_message["wav_name"] = wav_name
|
first_message["wav_format"] = wav_format
|
first_message["is_speaking"] = true
|
first_message["hotwords"] = hotword_msg
|
first_message["itn"] = use_itn
|
|
bytes, _ := json.Marshal(first_message)
|
message := string(bytes)
|
|
// fmt.Println(audio_bytes)
|
// fmt.Println(stride)
|
// fmt.Println(chunk_num)
|
// fmt.Println(message)
|
|
err := websocketConn.WriteMessage(websocket.TextMessage, []byte(message))
|
if err != nil {
|
log.Println("Failed to send the message:", err)
|
return
|
}
|
|
is_speaking := true
|
for i := 0; i < chunk_num; i++ {
|
beg := i * stride
|
var data []byte
|
if i == chunk_num-1 {
|
data = audio_bytes[beg:]
|
} else {
|
data = audio_bytes[beg : beg+stride]
|
}
|
|
err = websocketConn.WriteMessage(websocket.BinaryMessage, data)
|
if err != nil {
|
fmt.Println("Failed to send audio data:", err)
|
return
|
}
|
|
if i == chunk_num-1 {
|
is_speaking = false
|
endMsg := map[string]bool{"is_speaking": is_speaking}
|
endMsgBytes, err := json.Marshal(endMsg)
|
if err != nil {
|
fmt.Println("JSON serialization failed:", err)
|
return
|
}
|
err = websocketConn.WriteMessage(websocket.TextMessage, endMsgBytes)
|
if err != nil {
|
fmt.Println("Failed to send the termination message:", err)
|
return
|
}
|
}
|
var sleepDuration time.Duration
|
if args.mode == "offline" {
|
sleepDuration = time.Millisecond
|
} else {
|
fmt.Println("timesleep: Currently, only offline mode is supported.")
|
// sleepDuration = time.Duration(60*float64(args.chunk_size[1])/float64(args.chunk_interval)) * time.Millisecond
|
return
|
}
|
time.Sleep(sleepDuration)
|
}
|
|
if args.mode != "offline" {
|
fmt.Println("Currently, only offline mode is supported.")
|
return
|
}
|
if args.mode == "offline" {
|
for !offline_msg_done {
|
time.Sleep(1 * time.Second)
|
}
|
}
|
websocketConn.Close()
|
}
|
|
func message(id string) {
|
text_print := ""
|
var ibestWriter *os.File
|
var err error
|
if args.output_dir != "" {
|
filePath := fmt.Sprintf("%s/text.%d", args.output_dir, id)
|
ibestWriter, err = os.OpenFile(filePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
|
if err != nil {
|
log.Fatalf("failed to open file: %v", err)
|
}
|
} else {
|
ibestWriter = nil
|
}
|
|
for {
|
_, message, err := websocketConn.ReadMessage()
|
if err != nil {
|
log.Println("read:", err)
|
break
|
}
|
var meg Message
|
var wav_name string
|
timestamp := ""
|
|
err = json.Unmarshal(message, &meg)
|
if err != nil {
|
log.Println("unmarshal:", err)
|
continue
|
}
|
|
if meg.WavName != "" {
|
wav_name = meg.WavName
|
} else {
|
wav_name = "demo"
|
}
|
|
text := meg.Text
|
|
if meg.TimeStamp != "" {
|
timestamp = meg.TimeStamp
|
}
|
|
if ibestWriter != nil {
|
var text_write_line string
|
if timestamp != "" {
|
text_write_line = fmt.Sprintf("%s\t%s\t%s\n", wav_name, text, timestamp)
|
} else {
|
text_write_line = fmt.Sprintf("%s\t%s\n", wav_name, text)
|
}
|
_, err = ibestWriter.WriteString(text_write_line)
|
if err != nil {
|
log.Fatalf("Failed to write to file: %v", err)
|
}
|
}
|
|
if meg.Mode != "offline" {
|
fmt.Println("Currently, only offline mode is supported.")
|
return
|
}
|
|
if meg.Mode == "offline" {
|
if timestamp != "" {
|
text_print += fmt.Sprintf("%s timestamp: %s", text, timestamp)
|
} else {
|
text_print += fmt.Sprintf("%s ", text)
|
}
|
fmt.Println("\rpid" + id + ": " + wav_name + ": " + text_print)
|
offline_msg_done = true
|
}
|
}
|
}
|
|
func wsClient(id, chunk_begin, chunk_size int, done chan bool) {
|
for i := chunk_begin; i < chunk_begin+chunk_size; i++ {
|
offline_msg_done = false
|
|
u := url.URL{Scheme: "ws", Host: fmt.Sprintf("%s:%s", args.host, args.port), Path: "/"}
|
fmt.Printf("Thread %d: Connecting to %s\n", id, u.String())
|
var err error
|
websocketConn, _, err = websocket.DefaultDialer.Dial(u.String(), nil)
|
if err != nil {
|
log.Fatal("dial:", err)
|
}
|
defer websocketConn.Close()
|
|
var wg sync.WaitGroup
|
wg.Add(2)
|
|
go func() {
|
defer wg.Done()
|
recordFromScp(i, 1)
|
}()
|
|
go func() {
|
defer wg.Done()
|
id_str := strconv.Itoa(id)
|
i_str := strconv.Itoa(i)
|
message(id_str + "_" + i_str)
|
}()
|
|
wg.Wait()
|
}
|
|
done <- true
|
}
|
|
func oneThread(id, chunk_begin, chunk_size int, wg *sync.WaitGroup) {
|
defer wg.Done()
|
done := make(chan bool)
|
|
go wsClient(id, chunk_begin, chunk_size, done)
|
|
select {
|
case <-done:
|
fmt.Printf("Thread %d: Task completed\n", id)
|
}
|
}
|
|
func main() {
|
args.audio_in = "../audio/asr_example.wav"
|
args.thread_num = 1
|
args.host = "127.0.0.1"
|
args.port = "10095"
|
args.output_dir = "/workspace/models/Outputs"
|
args.hotword = "/workspace/models/hotword.txt"
|
args.chunk_size = []int{5, 10, 5}
|
args.chunk_interval = 10
|
args.mode = "offline"
|
args.audio_fs = 16000
|
args.use_itn = 1
|
var chunk_size, remain_wavs int
|
|
wavs := []string{args.audio_in}
|
|
total_len := len(wavs)
|
if total_len >= args.thread_num {
|
chunk_size = total_len / args.thread_num
|
remain_wavs = total_len - chunk_size*args.thread_num
|
} else {
|
chunk_size = 1
|
remain_wavs = 0
|
}
|
|
var wg sync.WaitGroup
|
chunk_begin := 0
|
for i := 0; i < args.thread_num; i++ {
|
wg.Add(1)
|
now_chunk_size := chunk_size
|
if remain_wavs > 0 {
|
now_chunk_size = chunk_size + 1
|
remain_wavs = remain_wavs - 1
|
}
|
go oneThread(i, chunk_begin, now_chunk_size, &wg)
|
|
chunk_begin = chunk_begin + now_chunk_size
|
}
|
wg.Wait()
|
}
|
