91 lines
2.8 KiB
Go
91 lines
2.8 KiB
Go
package handlers
|
||
|
||
import (
|
||
"encoding/binary"
|
||
"encoding/json"
|
||
"fmt"
|
||
"net"
|
||
"regexp"
|
||
"sync"
|
||
|
||
"rtp-app/pkg/rtp"
|
||
|
||
vosk "github.com/alphacep/vosk-api/go"
|
||
)
|
||
|
||
const (
|
||
Rate = 16000
|
||
bufSize = 2048
|
||
)
|
||
|
||
type VoskResult struct {
|
||
Text string `json:"text"`
|
||
}
|
||
|
||
func HandleConnection(conn *net.UDPConn, wg *sync.WaitGroup) {
|
||
model, err := vosk.NewModel("model-ru")
|
||
if err != nil {
|
||
fmt.Printf("Не удалось открыть модель Vosk: %v", err)
|
||
}
|
||
recognizer, err := vosk.NewRecognizer(model, float64(Rate))
|
||
if err != nil {
|
||
fmt.Printf("Не удалось запустить модель: %v", err)
|
||
}
|
||
defer wg.Done()
|
||
buffer := make([]byte, bufSize)
|
||
mu := sync.Mutex{}
|
||
re := regexp.MustCompile(`[\.\:]`)
|
||
for {
|
||
n, addr, err := conn.ReadFromUDP(buffer)
|
||
if err != nil {
|
||
fmt.Println("Ошибка при чтении данных:", err)
|
||
return
|
||
}
|
||
mu.Lock()
|
||
CallID := fmt.Sprintf("%s-%d", re.ReplaceAllString(addr.IP.String(), "-"), addr.Port)
|
||
mu.Unlock()
|
||
wg.Add(1)
|
||
processRTP(buffer[:n], CallID, wg, recognizer)
|
||
}
|
||
}
|
||
|
||
func processRTP(data []byte, CallID string, wg *sync.WaitGroup, recognizer *vosk.VoskRecognizer) {
|
||
packet := &rtp.Packet{}
|
||
if err := packet.Unmarshal(data); err != nil {
|
||
fmt.Printf("Ошибка при декодировании RTP пакета от %s: %v\n", CallID, err)
|
||
return
|
||
}
|
||
if recognizer.AcceptWaveform(BeToLe(packet.Payload)) != 0 {
|
||
resultStr := recognizer.Result()
|
||
var result VoskResult
|
||
if err := json.Unmarshal([]byte(resultStr), &result); err != nil {
|
||
fmt.Printf("Ошибка демаршалинга результата: %v", err)
|
||
}
|
||
if result.Text != "" {
|
||
fmt.Printf("Полная транскрипция: %s\t%s\n", CallID, result.Text)
|
||
}
|
||
} else {
|
||
partialResultStr := recognizer.PartialResult()
|
||
var partialResult VoskResult
|
||
if err := json.Unmarshal([]byte(partialResultStr), &partialResult); err != nil {
|
||
fmt.Printf("Ошибка демаршалинга частичного результата: %v", err)
|
||
return
|
||
}
|
||
if partialResult.Text != "" {
|
||
fmt.Printf("Частичная транскрипция: %s\t%s\n", CallID, partialResult.Text)
|
||
}
|
||
}
|
||
}
|
||
|
||
func BeToLe(input []byte) []byte {
|
||
if len(input)%4 != 0 {
|
||
panic("Длина входных данных должна быть кратна 4")
|
||
}
|
||
output := make([]byte, len(input))
|
||
for i := 0; i < len(input); i += 4 {
|
||
binary.LittleEndian.PutUint32(output[i:i+4], binary.BigEndian.Uint32(input[i:i+4]))
|
||
}
|
||
return output
|
||
}
|
||
|