Alexey Khit
161 lines
3.2 KiB
Go
161 lines
3.2 KiB
Go
// Package annexb - universal for H264 and H265
|
|
package annexb
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/binary"
|
|
)
|
|
|
|
const StartCode = "\x00\x00\x00\x01"
|
|
const startAUD = StartCode + "\x09\xF0"
|
|
const startAUDstart = startAUD + StartCode
|
|
|
|
// EncodeToAVCC
|
|
// will change original slice data!
|
|
// safeAppend should be used if original slice has useful data after end (part of other slice)
|
|
//
|
|
// FFmpeg MPEG-TS: 00000001 AUD 00000001 SPS 00000001 PPS 000001 IFrame
|
|
// FFmpeg H264: 00000001 SPS 00000001 PPS 000001 IFrame 00000001 PFrame
|
|
func EncodeToAVCC(b []byte, safeAppend bool) []byte {
|
|
const minSize = len(StartCode) + 1
|
|
|
|
// 1. Check frist "start code"
|
|
if len(b) < len(startAUDstart) || string(b[:len(StartCode)]) != StartCode {
|
|
return nil
|
|
}
|
|
|
|
// 2. Skip Access unit delimiter (AUD) from FFmpeg
|
|
if string(b[:len(startAUDstart)]) == startAUDstart {
|
|
b = b[6:]
|
|
}
|
|
|
|
var start int
|
|
|
|
for i, n := minSize, len(b)-minSize; i < n; {
|
|
// 3. Check "start code" (first 2 bytes)
|
|
if b[i] != 0 || b[i+1] != 0 {
|
|
i++
|
|
continue
|
|
}
|
|
|
|
// 4. Check "start code" (3 bytes size or 4 bytes size)
|
|
if b[i+2] == 1 {
|
|
if safeAppend {
|
|
// protect original slice from "damage"
|
|
b = bytes.Clone(b)
|
|
safeAppend = false
|
|
}
|
|
|
|
// convert start code from 3 bytes to 4 bytes
|
|
b = append(b, 0)
|
|
copy(b[i+1:], b[i:])
|
|
n++
|
|
} else if b[i+2] != 0 || b[i+3] != 1 {
|
|
i++
|
|
continue
|
|
}
|
|
|
|
// 5. Set size for previous AU
|
|
size := uint32(i - start - len(StartCode))
|
|
binary.BigEndian.PutUint32(b[start:], size)
|
|
|
|
start = i
|
|
|
|
i += minSize
|
|
}
|
|
|
|
// 6. Set size for last AU
|
|
size := uint32(len(b) - start - len(StartCode))
|
|
binary.BigEndian.PutUint32(b[start:], size)
|
|
|
|
return b
|
|
}
|
|
|
|
func DecodeAVCC(b []byte, safeClone bool) []byte {
|
|
if safeClone {
|
|
b = bytes.Clone(b)
|
|
}
|
|
for i := 0; i < len(b); {
|
|
size := int(binary.BigEndian.Uint32(b[i:]))
|
|
b[i] = 0
|
|
b[i+1] = 0
|
|
b[i+2] = 0
|
|
b[i+3] = 1
|
|
i += 4 + size
|
|
}
|
|
return b
|
|
}
|
|
|
|
// DecodeAVCCWithAUD - AUD doesn't important for FFmpeg, but important for Safari
|
|
func DecodeAVCCWithAUD(src []byte) []byte {
|
|
dst := make([]byte, len(startAUD)+len(src))
|
|
copy(dst, startAUD)
|
|
copy(dst[len(startAUD):], src)
|
|
DecodeAVCC(dst[len(startAUD):], false)
|
|
return dst
|
|
}
|
|
|
|
const (
|
|
h264PFrame = 1
|
|
h264IFrame = 5
|
|
h264SPS = 7
|
|
h264PPS = 8
|
|
|
|
h265VPS = 32
|
|
h265PFrame = 1
|
|
)
|
|
|
|
// IndexFrame - get new frame start position in the AnnexB stream
|
|
func IndexFrame(b []byte) int {
|
|
if len(b) < len(startAUDstart) {
|
|
return -1
|
|
}
|
|
|
|
for i := len(startAUDstart); ; {
|
|
if di := bytes.Index(b[i:], []byte(StartCode)); di < 0 {
|
|
break
|
|
} else {
|
|
i += di + 4 // move to NALU start
|
|
}
|
|
|
|
if i >= len(b) {
|
|
break
|
|
}
|
|
|
|
h264Type := b[i] & 0b1_1111
|
|
switch h264Type {
|
|
case h264PFrame, h264SPS:
|
|
return i - 4 // move to start code
|
|
case h264IFrame, h264PPS:
|
|
continue
|
|
}
|
|
|
|
h265Type := (b[i] >> 1) & 0b11_1111
|
|
switch h265Type {
|
|
case h265PFrame, h265VPS:
|
|
return i - 4 // move to start code
|
|
}
|
|
}
|
|
|
|
return -1
|
|
}
|
|
|
|
func FixAnnexBInAVCC(b []byte) []byte {
|
|
for i := 0; i < len(b); {
|
|
if i+4 >= len(b) {
|
|
break
|
|
}
|
|
|
|
size := bytes.Index(b[i+4:], []byte{0, 0, 0, 1})
|
|
if size < 0 {
|
|
size = len(b) - (i + 4)
|
|
}
|
|
|
|
binary.BigEndian.PutUint32(b[i:], uint32(size))
|
|
|
|
i += size + 4
|
|
}
|
|
|
|
return b
|
|
}
|