vosk-recognize-dev/pkg/wav/encoder.go

package wav

import (
    "bytes"
    "encoding/binary"
    "errors"
    "fmt"
    "io"
    "os"
    "time"

    "rtp-app/pkg/audio"
    "rtp-app/pkg/riff"
)

type Encoder struct {
    w               io.WriteSeeker
    buf             *bytes.Buffer
    SampleRate      int
    BitDepth        int
    NumChans        int
    WavAudioFormat  int
    Metadata        *Metadata
    WrittenBytes    int
    frames          int
    pcmChunkStarted bool
    pcmChunkSizePos int
    wroteHeader     bool
}

func NewEncoder(w io.WriteSeeker, sampleRate, bitDepth, numChans, audioFormat int) *Encoder {
    return &Encoder{
        w:              w,
        buf:            bytes.NewBuffer(make([]byte, 0, bytesNumFromDuration(time.Minute, sampleRate, bitDepth)*numChans)),
        SampleRate:     sampleRate,
        BitDepth:       bitDepth,
        NumChans:       numChans,
        WavAudioFormat: audioFormat,
    }
}

func (e *Encoder) AddLE(src interface{}) error {
    e.WrittenBytes += binary.Size(src)
    return binary.Write(e.w, binary.LittleEndian, src)
}

func (e *Encoder) AddBE(src interface{}) error {
    e.WrittenBytes += binary.Size(src)
    return binary.Write(e.w, binary.LittleEndian, src)
}

func (e *Encoder) addBuffer(buf *audio.IntBuffer) error {
    if buf == nil {
        return fmt.Errorf("can't add a nil buffer")
    }
    frameCount := buf.NumFrames()
    var err error
    for i := 0; i < frameCount; i++ {
        for j := 0; j < buf.Format.NumChannels; j++ {
            v := buf.Data[i*buf.Format.NumChannels+j]
            switch e.BitDepth {
                case 8:
                    if err = binary.Write(e.buf, binary.LittleEndian, uint8(v)); err != nil {
                        return err
                    }
                case 16:
                    if err = binary.Write(e.buf, binary.LittleEndian, int16(v)); err != nil {
                        return err
                    }
                case 24:
                    if err = binary.Write(e.buf, binary.LittleEndian, audio.Int32toInt24LEBytes(int32(v))); err != nil {
                        return err
                    }
                case 32:
                    if err = binary.Write(e.buf, binary.LittleEndian, int32(v)); err != nil {
                        return err
                    }
                default:
                    return fmt.Errorf("can't add frames of bit size %d", e.BitDepth)
            }
        }
        e.frames++
    }
    if n, err := e.w.Write(e.buf.Bytes()); err != nil {
        e.WrittenBytes += n
        return err
    }
    e.WrittenBytes += e.buf.Len()
    e.buf.Reset()
    return nil
}

func (e *Encoder) writeHeader() error {
    if e.wroteHeader {
        return errors.New("заголовок уже присутствует")
    }
    e.wroteHeader = true
    if e == nil {
        return fmt.Errorf("не могу записать encoder=nil")
    }
    if e.w == nil {
        return fmt.Errorf("не могу записать writer=nil")
    }

    if e.WrittenBytes > 0 {
        return nil
    }
    if err := e.AddLE(riff.RiffID); err != nil {
        return err
    }
    if err := e.AddLE(uint32(42)); err != nil {
        return err
    }
    if err := e.AddLE(riff.WavFormatID); err != nil {
        return err
    }
    if err := e.AddLE(riff.FmtID); err != nil {
        return err
    }
    if err := e.AddLE(uint32(16)); err != nil {
        return err
    }
    if err := e.AddLE(uint16(e.WavAudioFormat)); err != nil {
        return err
    }
    if err := e.AddLE(uint16(e.NumChans)); err != nil {
        return fmt.Errorf("ошибка кодирования количества каналов - %v", err)
    }
    if err := e.AddLE(uint32(e.SampleRate)); err != nil {
        return fmt.Errorf("ошибка частоты дискретизации - %v", err)
    }
    blockAlign := e.NumChans * e.BitDepth / 8
    if err := e.AddLE(uint32(e.SampleRate * blockAlign)); err != nil {
        return fmt.Errorf("ошибка среднего значения битрейта - %v", err)
    }
    if err := e.AddLE(uint16(blockAlign)); err != nil {
        return err
    }
    if err := e.AddLE(uint16(e.BitDepth)); err != nil {
        return fmt.Errorf("ошибка количества бит на выборку - %v", err)
    }
    return nil
}

func (e *Encoder) Write(buf *audio.IntBuffer) error {
    if !e.wroteHeader {
        if err := e.writeHeader(); err != nil {
            return err
        }
    }
    if !e.pcmChunkStarted {
        if err := e.AddLE(riff.DataFormatID); err != nil {
            return fmt.Errorf("ошибка кодирования заголовка %v", err)
        }
        e.pcmChunkStarted = true
        e.pcmChunkSizePos = e.WrittenBytes
        if err := e.AddLE(uint32(42)); err != nil {
            return fmt.Errorf("%v размер заголовка при записи фрагмента данных wav", err)
        }
    }
    return e.addBuffer(buf)
}

func (e *Encoder) WriteFrame(value interface{}) error {
    if !e.wroteHeader {
        e.writeHeader()
    }
    if !e.pcmChunkStarted {
        if err := e.AddLE(riff.DataFormatID); err != nil {
            return fmt.Errorf("error encoding sound header %v", err)
        }
        e.pcmChunkStarted = true
        e.pcmChunkSizePos = e.WrittenBytes
        if err := e.AddLE(uint32(42)); err != nil {
            return fmt.Errorf("%v when writing wav data chunk size header", err)
        }
    }
    e.frames++
    return e.AddLE(value)
}

func (e *Encoder) writeMetadata() error {
    chunkData := encodeInfoChunk(e)
    if err := e.AddBE(CIDList); err != nil {
        return fmt.Errorf("не удалось записать перечень chunk ID: %s", err)
    }
    if err := e.AddLE(uint32(len(chunkData))); err != nil {
        return fmt.Errorf("не удалось записать перечень chunk size: %s", err)
    }
    return e.AddBE(chunkData)
}

func (e *Encoder) Close() error {
    if e == nil || e.w == nil {
        return nil
    }
    if e.Metadata != nil {
        if err := e.writeMetadata(); err != nil {
            return fmt.Errorf("не удалось записать метаданные - %v", err)
        }
    }
    if _, err := e.w.Seek(4, 0); err != nil {
        return err
    }
    if err := e.AddLE(uint32(e.WrittenBytes) - 8); err != nil {
        return fmt.Errorf("%v общее количество записанных байтов", err)
    }
    if e.pcmChunkSizePos > 0 {
        if _, err := e.w.Seek(int64(e.pcmChunkSizePos), 0); err != nil {
            return err
        }
        chunksize := uint32((int(e.BitDepth) / 8) * int(e.NumChans) * e.frames)
        if err := e.AddLE(uint32(chunksize)); err != nil {
            return fmt.Errorf("%v размера заголовка при записи wav-данных", err)
        }
    }
    if _, err := e.w.Seek(0, 2); err != nil {
        return err
    }
    switch e.w.(type) {
        case *os.File:
            return e.w.(*os.File).Sync()
    }
    return nil
}