Merge pull request #34 from edgeware/annex-b

Added Annex B byte stream functions

README.md

@@ -141,3 +141,12 @@ Some code in pkg/mp4, comes from or is based on https://github.com/jfbus/mp4 whi
 
 Some code in pkg/bits comes from or is based on https://github.com/tcnksm/go-casper/tree/master/internal/bits
 `Copyright (c) 2017 Taichi Nakashima`.
+
+## Versions
+
+|Version | Highlight |
+|---|---|---|---|---|
+|0.14| Added functions to use Annex B byte stream for AVC/H.264|
+|0.13| Removed third-party log library |
+|0.12| Complete parsing of AVC/H.264 SPS and PPS |
+

avc/annexb.go

@@ -0,0 +1,111 @@
+package avc
+
+//Functions to handle AnnexB Byte stream format"
+
+import "encoding/binary"
+
+// ExtractNalusFromByteStream - extract NALUs without startcode from ByteStream
+func ExtractNalusFromByteStream(data []byte) [][]byte {
+	currNalStart := -1
+	n := len(data)
+	var nalus [][]byte
+	for i := 0; i < n-3; i++ {
+		if data[i] == 0 && data[i+1] == 0 && data[i+2] == 1 {
+			if currNalStart > 0 {
+				currNalEnd := i
+				for j := i - 1; j > currNalStart; j-- {
+					// Remove zeros from end of NAL unit
+					if data[j] == 0 {
+						currNalEnd = j
+					} else {
+						break
+					}
+				}
+				nalus = append(nalus, extractSlice(data, currNalStart, currNalEnd))
+			}
+			currNalStart = i + 3
+		}
+	}
+	if currNalStart < 0 {
+		return nil
+	}
+	nalus = append(nalus, extractSlice(data, currNalStart, n))
+	return nalus
+}
+
+func extractSlice(data []byte, start, stop int) []byte {
+	sl := make([]byte, stop-start)
+	_ = copy(sl, data[start:stop])
+	return sl
+}
+
+type scNalu struct {
+	startCodeLength int
+	startPos        int
+}
+
+// ConvertByteStreamToNaluSample - Change start codes to 4-byte length fields
+func ConvertByteStreamToNaluSample(stream []byte) []byte {
+	streamLen := len(stream)
+	var scNalus []scNalu
+	minStartCodeLength := 4
+	for i := 0; i < streamLen-3; i++ {
+		if stream[i] == 0 && stream[i+1] == 0 && stream[i+2] == 1 {
+			startCodeLength := 3
+			startPos := i + 3
+			if i >= 1 && stream[i-1] == 0 {
+				startCodeLength++
+			}
+			if startCodeLength < minStartCodeLength {
+				minStartCodeLength = startCodeLength
+			}
+			scNalus = append(scNalus, scNalu{startCodeLength, startPos})
+		}
+	}
+	lengthField := make([]byte, 4)
+	var nalLength int
+	if minStartCodeLength == 4 {
+		// In-place replacement of startcodes for length fields
+		for i, s := range scNalus {
+
+			if i+1 < len(scNalus) {
+				nalLength = scNalus[i+1].startPos - s.startPos - 4
+			} else {
+				nalLength = len(stream) - scNalus[i].startPos
+			}
+			binary.BigEndian.PutUint32(lengthField, uint32(nalLength))
+			copy(stream[s.startPos-4:s.startPos], lengthField)
+		}
+		return stream
+	} else {
+		// Build new output slice with one extra byte per NALU
+		out := make([]byte, 0, streamLen+len(scNalus))
+		for i, s := range scNalus {
+			if i+1 < len(scNalus) {
+				nalLength = scNalus[i+1].startPos - s.startPos - scNalus[i+1].startCodeLength
+			} else {
+				nalLength = len(stream) - scNalus[i].startPos
+			}
+			binary.BigEndian.PutUint32(lengthField, uint32(nalLength))
+			out = append(out, lengthField...)
+			out = append(out, stream[s.startPos:s.startPos+nalLength]...)
+		}
+		return out
+	}
+}
+
+// ConvertSampleToByteStream - Replace 4-byte NALU lengths with start codes
+func ConvertSampleToByteStream(sample []byte) []byte {
+	sampleLength := uint32(len(sample))
+	var pos uint32 = 0
+	for {
+		if pos >= sampleLength {
+			break
+		}
+		naluLength := binary.BigEndian.Uint32(sample[pos : pos+4])
+		startCode := []byte{0, 0, 0, 1}
+		copy(sample[pos:pos+4], startCode)
+		pos += naluLength + 4
+	}
+	return sample
+}

avc/annexb_test.go

@@ -0,0 +1,110 @@
+package avc
+
+import (
+	"testing"
+
+	"github.com/go-test/deep"
+)
+
+func TestNaluExtraction(t *testing.T) {
+	testCases := []struct {
+		name   string
+		input  []byte
+		wanted [][]byte
+	}{
+		{
+			"One 4-byte start-code NALU",
+			[]byte{0, 0, 0, 1, 2},
+			[][]byte{{2}},
+		},
+		{
+			"One 3-byte start-code NALU",
+			[]byte{0, 0, 1, 2},
+			[][]byte{{2}},
+		},
+		{
+			"No start-code",
+			[]byte{0, 0, 2},
+			nil,
+		},
+		{
+			"Just a start-code",
+			[]byte{0, 0, 1},
+			nil,
+		},
+		{
+			"Two NALUs (start codes)",
+			[]byte{0, 0, 1, 2, 0, 0, 0, 1, 1},
+			[][]byte{{2}, {1}},
+		},
+	}
+
+	for _, tc := range testCases {
+		got := ExtractNalusFromByteStream(tc.input)
+		if diff := deep.Equal(got, tc.wanted); diff != nil {
+			t.Errorf("%s: %v", tc.name, diff)
+		}
+	}
+}
+
+func TestByteStreamToNaluSampleConversion(t *testing.T) {
+	testCases := []struct {
+		name   string
+		input  []byte
+		wanted []byte
+	}{
+		{
+			"One 4-byte start-code + 2-byte NALU",
+			[]byte{0, 0, 0, 1, 2, 3},
+			[]byte{0, 0, 0, 2, 2, 3},
+		},
+		{
+			"One 3-byte start-code + 2-byte NALU",
+			[]byte{0, 0, 1, 2, 3},
+			[]byte{0, 0, 0, 2, 2, 3},
+		},
+		{
+			"Two 4-byte start-codes",
+			[]byte{0, 0, 0, 1, 2, 3, 0, 0, 0, 1, 7},
+			[]byte{0, 0, 0, 2, 2, 3, 0, 0, 0, 1, 7},
+		},
+		{
+			"Two 3-byte start-codes",
+			[]byte{0, 0, 1, 2, 3, 0, 0, 1, 7},
+			[]byte{0, 0, 0, 2, 2, 3, 0, 0, 0, 1, 7},
+		},
+	}
+
+	for _, tc := range testCases {
+		got := ConvertByteStreamToNaluSample(tc.input)
+		if diff := deep.Equal(got, tc.wanted); diff != nil {
+			t.Errorf("%s: %v", tc.name, diff)
+		}
+	}
+}
+
+func TestSampleToByteStreamConversion(t *testing.T) {
+	testCases := []struct {
+		name   string
+		input  []byte
+		wanted []byte
+	}{
+		{
+			"One NALU",
+			[]byte{0, 0, 0, 2, 2, 3},
+			[]byte{0, 0, 0, 1, 2, 3},
+		},
+		{
+			"Two NALUs",
+			[]byte{0, 0, 0, 2, 2, 3, 0, 0, 0, 1, 7},
+			[]byte{0, 0, 0, 1, 2, 3, 0, 0, 0, 1, 7},
+		},
+	}
+
+	for _, tc := range testCases {
+		got := ConvertSampleToByteStream(tc.input)
+		if diff := deep.Equal(got, tc.wanted); diff != nil {
+			t.Errorf("%s: %v", tc.name, diff)
+		}
+	}
+}