pion h264

main.go

package main

import (
    "io"
    "log"
    "os"
    "strconv"
    "time"
)

type FindNalState struct {
    PrefixCount int
    LastNullCount int
    buf []byte
}
func NewFindNalState() FindNalState {
    return FindNalState{PrefixCount: 0, LastNullCount: 0, buf: make([]byte, 0)}
}
func (h *FindNalState) NalScan(buf []byte) [][]byte {
    if len(h.buf) > 1024 * 1024 { panic("FindNalState buf len panic") }
    nals := make([][]byte, 0)

    // смещение в массиве buf после комбинации 0x00_00_01 или 0x00_00_00_01, указывает на следующий за ней байт
    var lastPrefixOffset *int = nil
    i := 0 // следующий байт в обработке
    for {
        if i >= len(buf) {
            // буффер закончился
            if lastPrefixOffset != nil {
                // копируем остаток с конца последнего префикса в буффер
                h.buf = make([]byte, 0)
                h.buf = append(h.buf, buf[*lastPrefixOffset:]...)
            } else {
                // префикса в данной пачке найдено не было, отправляем все данные в буффер
                h.buf = append(h.buf, buf...)
            }
            break
        }
        b := buf[i]; i += 1
        switch b {
        case 0x00: { if h.LastNullCount < 3 { h.LastNullCount += 1 }; continue }
        case 0x01: {
            if h.LastNullCount >= 2 {
                // нашли префикс 0x00_00_01 или 0x00_00_00_01
                prefixOffset := i
                if lastPrefixOffset != nil {
                    // отправляем NAL с последнего префикса до текущего
                    size := (i - h.LastNullCount) - *lastPrefixOffset - 1
                    if size > 0 && h.PrefixCount > 0 {
                        nal := buf[*lastPrefixOffset : *lastPrefixOffset + size]
                        // send_nal
                        nals = append(nals, nal)
                    }
                } else {
                    // префикса в данной пачке еще найдено не было, отправляем данные буффера и данные до текущего префикса
                    size := i - h.LastNullCount - 1

                    nal := make([]byte, 0)
                    if size < 0 {
                        if len(h.buf) > 0 {
                            nal = append(nal, h.buf[0 : len(h.buf) + size]...)
                        }
                    } else {
                        nal = append(nal, h.buf...)
                        nal = append(nal, buf[0 : size]...)
                    }
                    // остылаем только непустой NAL и только после обнаруженного префикса
                    if len(nal) > 0 && h.PrefixCount > 0 {
                        // send_nal
                        nals = append(nals, nal)
                    }
                    h.buf = make([]byte, 0)
                }
                p := prefixOffset
                lastPrefixOffset = &p
                h.PrefixCount += 1
            }
        }
        default:
        }
        h.LastNullCount = 0
    }
    return nals
}

type NalUnitType uint8
const( //   Table 7-1 NAL unit type codes
    Unspecified NalUnitType = 0                // Unspecified
    CodedSliceNonIdr NalUnitType = 1           // Coded slice of a non-IDR picture
    CodedSliceDataPartitionA NalUnitType = 2   // Coded slice data partition A
    CodedSliceDataPartitionB NalUnitType = 3   // Coded slice data partition B
    CodedSliceDataPartitionC NalUnitType = 4   // Coded slice data partition C
    CodedSliceIdr NalUnitType = 5              // Coded slice of an IDR picture
    SEI NalUnitType = 6                        // Supplemental enhancement information (SEI)
    SPS NalUnitType = 7                        // Sequence parameter set
    PPS NalUnitType = 8                        // Picture parameter set
    AUD NalUnitType = 9                        // Access unit delimiter
    EndOfSequence NalUnitType = 10             // End of sequence
    EndOfStream NalUnitType = 11               // End of stream
    Filler NalUnitType = 12                    // Filler data
    SpsExt NalUnitType = 13                    // Sequence parameter set extension
    // 14..18           // Reserved
    NalUnitTypeCodedSliceAux NalUnitType = 19  // Coded slice of an auxiliary coded picture without partitioning
    // 20..23           // Reserved
    // 24..31           // Unspecified
)
func NalUnitTypeStr(v NalUnitType) string {
    str := "Unknown"
    switch v {
    case 0: { str = "Unspecified" }
    case 1: { str = "CodedSliceNonIdr" }
    case 2: { str = "CodedSliceDataPartitionA" }
    case 3: { str = "CodedSliceDataPartitionB" }
    case 4: { str = "CodedSliceDataPartitionC" }
    case 5: { str = "CodedSliceIdr" }
    case 6: { str = "SEI" }
    case 7: { str = "SPS" }
    case 8: { str = "PPS" }
    case 9: { str = "AUD" }
    case 10: { str = "EndOfSequence" }
    case 11: { str = "EndOfStream" }
    case 12: { str = "Filler" }
    case 13: { str = "SpsExt" }
    case 19: { str = "NalUnitTypeCodedSliceAux" }
    default: { str = "Unknown" }
    }
    str = str + "(" + strconv.FormatInt(int64(v), 10) + ")"
    return str
}

type Nal struct {
    PictureOrderCount uint32

    // NAL header
    ForbiddenZeroBit bool
    RefIdc uint8
    UnitType NalUnitType

    HeaderByte uint8
    Rbsp []byte

    Data []byte // header_byte + rbsp
}
func NewNal() Nal {
    return Nal {PictureOrderCount: 0, ForbiddenZeroBit: false, RefIdc: 0, UnitType: Unspecified, HeaderByte: 0, Rbsp: make([]byte, 0), Data: make([]byte, 0)}
}
func (h *Nal) ParseHeader(firstByte byte) {
    h.HeaderByte = firstByte
    h.ForbiddenZeroBit = (((firstByte & 0b10000000) >> 7) == 1)
    h.RefIdc =             (firstByte & 0b01100000) >> 5
    h.UnitType = NalUnitType((firstByte & 0b00011111) >> 0)
}

func loadFile(nalChn chan Nal) {
    file, err := os.Open("/Users/user/dev/rust_wasm/output_1920x1080_baseline_tag.h264")
    if err != nil { panic(err) }

    nals := make([][]byte, 0)
    nalStream := NewFindNalState()
    for {
        buf := make([]byte, 1024)
        n, err := file.Read(buf)
        if err != nil && err != io.EOF { log.Fatal("Error Reading: ", err); break }
        if n == 0 { break }
        nal := nalStream.NalScan(buf[0:n])
        nals = append(nals, nal...)
    }
    log.Println("nals: ", len(nals))

    i := 0
    pic := int64(0)
    start := time.Now()
    for {
        if i >= len(nals) { i = 0 }
        nalData := nals[i]; i+=1
        nal := NewNal()
        nal.HeaderByte = nalData[0]
        nal.ParseHeader(nal.HeaderByte)
        if nal.UnitType == SEI { continue }
        nal.Data = nalData
        nal.Rbsp = nalData[1:]
        // log.Println(i, "nal: ", NalUnitTypeStr(nal.UnitType))
        if nal.UnitType == CodedSliceNonIdr || nal.UnitType == CodedSliceIdr {
            pic+=1
            currentMustBeUs := int64(float64(pic) * 1000000.0 / 25.0)
            now := time.Now()
            fromStart := now.Sub(start).Microseconds()
            delta := currentMustBeUs - fromStart
            if delta > 0 {
                time.Sleep(time.Duration(delta) * time.Microsecond)
            }
            //if fromStart > 0 {
            //    log.Println("fps: ", float64(pic) / (float64(fromStart) / 1000000.0))
            //}
        }
        nalChn <- nal
    }
}






func main3() {
    var nalStream = make(chan Nal)
    go loadFile(nalStream)

    // example
    // go start_webrtc(nalStream)
    select { }
}

// webrtc peer with videoTrack send nals from
//func (p Peer) SendH264() {
//    i := 0
//    for {
//        nal := <- nalStream
//        Samples := uint32(0)
//        if nal.UnitType == CodedSliceNonIdr || nal.UnitType == CodedSliceIdr {
//            //Samples = 90000
//            Samples = 90000
//        }
//        // try to add prefix 0x00_00_00_01
//        // sample := media.Sample{Data: append([]byte{0x00, 0x00, 0x00, 0x01}, nal.Data...), Samples: Samples}
//
//        // without prefix
//        sample := media.Sample{Data: nal.Data, Samples: Samples}
//
//        if nal.UnitType == CodedSliceNonIdr ||
//            nal.UnitType == CodedSliceIdr ||
//            nal.UnitType == SPS || nal.UnitType == PPS {
//            log.Println(i, "send nal ", NalUnitTypeStr(nal.UnitType))
//            err := p.videoTrack.WriteSample(sample)
//            if err != nil { panic(err) }
//            i++
//        }
//    }
//}