Merge pull request #36 from edgeware/avc-improvements

AVC improvements

Author: tobbee

README.md

@@ -144,9 +144,4 @@ Some code in pkg/bits comes from or is based on https://github.com/tcnksm/go-cas
 
 ## Versions
 
-| Version | Highlight |
-| ------  | --------- |
-| 0.14.0 | Added functions to use Annex B byte stream for AVC/H.264 |
-| 0.13.0 | Removed third-party log library |
-| 0.12.0 | Complete parsing of AVC/H.264 SPS and PPS |
-
+See [Versions.md](Versions.md).

Versions.md

@@ -0,0 +1,9 @@
+# Versions
+
+| Version | Highlight |
+| ------  | --------- |
+
+| 0.15.0 | Support for multiple SPS and more explicit NALU types for AVC |
+| 0.14.0 | Added functions to use Annex B byte stream for AVC/H.264 |
+| 0.13.0 | Removed third-party log library |
+| 0.12.0 | Complete parsing of AVC/H.264 SPS and PPS |

avc/avc.go

@@ -2,13 +2,16 @@ package avc
 
 import (
 	"encoding/binary"
+	"fmt"
 )
 
 // NalType - AVC nal type
 type NalType uint16
 
 const (
-	// NALU_IDR - IDR Random Access Picture NAL Unit
+	// NALU_NON_IDR - Non-IDR Slice NAL unit
+	NALU_NON_IDR = NalType(1)
+	// NALU_IDR - IDR Random Access Slice NAL Unit
 	NALU_IDR = NalType(5)
 	// NALU_SEI - Supplementary Enhancement Information NAL Unit
 	NALU_SEI = NalType(6)
@@ -18,10 +21,12 @@ const (
 	NALU_PPS = NalType(8)
 	// NALU_AUD - AccessUnitDelimiter NAL Unit
 	NALU_AUD = NalType(9)
+	// NALU_EO_SEQ - End of Sequence NAL Unit
+	NALU_EO_SEQ = NalType(10)
+	// NALU_EO_STREAM - End of Stream NAL Unit
+	NALU_EO_STREAM = NalType(11)
 	// NALU_FILL - Filler NAL Unit
 	NALU_FILL = NalType(12)
-	// ExtendedSAR - Extended Sample Aspect Ratio Code
-	ExtendedSAR = 255
 )
 
 func (a NalType) String() string {
@@ -37,7 +42,7 @@ func (a NalType) String() string {
 	case NALU_AUD:
 		return "AUD"
 	default:
-		return "other"
+		return fmt.Sprintf("Other %d", a)
 	}
 }
 
@@ -103,3 +108,25 @@ func HasParameterSets(b []byte) bool {
 	}
 	return false
 }
+
+// GetParameterSets - get (multiple) SPS and PPS from a sample
+func GetParameterSets(sample []byte) (sps [][]byte, pps [][]byte) {
+	sampleLength := uint32(len(sample))
+	var pos uint32 = 0
+	for {
+		if pos >= sampleLength {
+			break
+		}
+		nalLength := binary.BigEndian.Uint32(sample[pos : pos+4])
+		pos += 4
+		nalHdr := sample[pos]
+		switch GetNalType(nalHdr) {
+		case NALU_SPS:
+			sps = append(sps, sample[pos:pos+nalLength])
+		case NALU_PPS:
+			pps = append(pps, sample[pos:pos+nalLength])
+		}
+		pos += nalLength
+	}
+	return sps, pps
+}

avc/avcdecoderconfigurationrecord.go

@@ -26,8 +26,9 @@ type AVCDecConfRec struct {
 }
 
 // CreateAVCDecConfRec - Create an AVCDecConfRec based on SPS and PPS
-func CreateAVCDecConfRec(spsNALU []byte, ppsNALUs [][]byte) AVCDecConfRec {
-	sps, err := ParseSPSNALUnit(spsNALU, false) // false -> parse only start of VUI
+func CreateAVCDecConfRec(spsNALUs [][]byte, ppsNALUs [][]byte) AVCDecConfRec {
+
+	sps, err := ParseSPSNALUnit(spsNALUs[0], false) // false -> parse only start of VUI
 	if err != nil {
 		panic("Could not parse SPS")
 	}
@@ -36,7 +37,7 @@ func CreateAVCDecConfRec(spsNALU []byte, ppsNALUs [][]byte) AVCDecConfRec {
 		AVCProfileIndication: byte(sps.Profile),
 		ProfileCompatibility: byte(sps.ProfileCompatibility),
 		AVCLevelIndication:   byte(sps.Level),
-		SPSnalus:             [][]byte{spsNALU},
+		SPSnalus:             spsNALUs,
 		PPSnalus:             ppsNALUs,
 		ChromaFormat:         1,
 		BitDepthLumaMinus1:   0,

avc/sps.go

@@ -8,6 +8,9 @@ import (
 	"github.com/edgeware/mp4ff/bits"
 )
 
+// extendedSAR - Extended Sample Aspect Ratio Code
+const extendedSAR = 255
+
 var ErrNotSPS = errors.New("Not an SPS NAL unit")
 
 // SPS - AVC SPS parameters
@@ -355,7 +358,7 @@ func parseVUI(reader *bits.EBSPReader, parseVUIBeyondAspectRatio bool) (*VUIPara
 		if err != nil {
 			return nil, err
 		}
-		if aspectRatioIDC == ExtendedSAR {
+		if aspectRatioIDC == extendedSAR {
 			vui.SampleAspectRatioWidth, err = reader.Read(16)
 			if err != nil {
 				return nil, err

bits/bits.go

@@ -80,7 +80,7 @@ func NewReader(rd io.Reader) *Reader {
 func (r *Reader) Read(n int) (uint, error) {
 	var err error
 
-	for r.n <= n {
+	for r.n < n {
 		r.v <<= 8
 		var b uint8
 		err = binary.Read(r.rd, binary.BigEndian, &b)
@@ -112,7 +112,7 @@ func (r *Reader) ReadFlag() (bool, error) {
 func (r *Reader) MustRead(n int) uint {
 	var err error
 
-	for r.n <= n {
+	for r.n < n {
 		r.v <<= 8
 		var b uint8
 		err = binary.Read(r.rd, binary.BigEndian, &b)

bits/ebsp.go

@@ -50,12 +50,11 @@ func (r *EBSPReader) MustRead(n int) uint {
 			}
 			r.pos++
 			r.zeroCount = 0
+		}
+		if b != 0 {
+			r.zeroCount = 0
 		} else {
-			if b != 0 {
-				r.zeroCount = 0
-			} else {
-				r.zeroCount++
-			}
+			r.zeroCount++
 		}
 		r.v |= uint(b)
 
@@ -151,12 +150,11 @@ func (r *EBSPReader) Read(n int) (uint, error) {
 			}
 			r.pos++
 			r.zeroCount = 0
+		}
+		if b != 0 {
+			r.zeroCount = 0
 		} else {
-			if b != 0 {
-				r.zeroCount = 0
-			} else {
-				r.zeroCount++
-			}
+			r.zeroCount++
 		}
 		r.v |= uint(b)
 

bits/esbp_test.go bits/ebsp_test.gobits/esbp_test.go renamed to bits/ebsp_test.go

@@ -70,6 +70,11 @@ func TestEbspParser(t *testing.T) {
 			"27640020ac2ec05005bb011000000300100000078e840016e300005b8d8bdef83b438627",
 			"27640020ac2ec05005bb0110000000100000078e840016e300005b8d8bdef83b438627",
 		},
+		{
+			"Long zero sequence",
+			"00000300000300",
+			"0000000000",
+		},
 	}
 
 	for _, c := range cases {
@@ -212,3 +217,34 @@ func TestReadTrailingRbspBits(t *testing.T) {
 		t.Errorf("Not at end after reading rbsp_trailing_bits")
 	}
 }
+
+func TestEBSPWriter(t *testing.T) {
+	testCases := []struct {
+		in  []byte
+		out []byte
+	}{
+		{
+			in:  []byte{0, 0, 0, 1},
+			out: []byte{0, 0, 3, 0, 1},
+		},
+		{
+			in:  []byte{1, 0, 0, 2},
+			out: []byte{1, 0, 0, 3, 2},
+		},
+		{
+			in:  []byte{0, 0, 0, 0, 0},
+			out: []byte{0, 0, 3, 0, 0, 3, 0},
+		},
+	}
+	for _, tc := range testCases {
+		buf := bytes.Buffer{}
+		w := NewEBSPWriter(&buf)
+		for _, b := range tc.in {
+			w.Write(uint(b), 8)
+		}
+		diff := deep.Equal(buf.Bytes(), tc.out)
+		if diff != nil {
+			t.Errorf("Got %v but wanted %d", buf.Bytes(), tc.out)
+		}
+	}
+}

bits/ebsp_writer.go

@@ -0,0 +1,68 @@
+package bits
+
+import (
+	"encoding/binary"
+	"io"
+)
+
+// EbspWriter writes bits into underlying io.Writer and inserts
+// start-code emulation prevention bytes as necessary
+// Stops writing at first error.
+// Errors that have occured can later be checked with Error().
+type EBSPWriter struct {
+	n   int   // current number of bits
+	v   uint  // current accumulated value
+	err error // The first error caused by any write operation
+	nr0 int   // Number preceeding zero bytes
+
+	wr io.Writer
+}
+
+// NewWriter - returns a new Writer
+func NewEBSPWriter(w io.Writer) *EBSPWriter {
+	return &EBSPWriter{
+		wr: w,
+	}
+}
+
+// Write - write n bits from bits and save error state
+func (w *EBSPWriter) Write(bits uint, n int) {
+	if w.err != nil {
+		return
+	}
+	w.v <<= uint(n)
+	w.v |= bits & mask(n)
+	w.n += n
+	for w.n >= 8 {
+		b := (w.v >> (uint(w.n) - 8)) & mask(8)
+		if w.nr0 == 2 && b <= 3 {
+			if err := binary.Write(w.wr, binary.BigEndian, uint8(3)); err != nil {
+				w.err = err
+				return
+			}
+			w.nr0 = 0
+		}
+		if err := binary.Write(w.wr, binary.BigEndian, uint8(b)); err != nil {
+			w.err = err
+			return
+		}
+		if b == 0 {
+			w.nr0++
+		}
+		w.n -= 8
+	}
+	w.v &= mask(8)
+}
+
+// Write - write rbsp trailing bits (a 1 followed by zeros to a byte boundary)
+func (w *EBSPWriter) WriteRbspTrailingBits() {
+	w.Write(1, 1)
+	w.StuffByteWithZeros()
+}
+
+// StuffByteWithZeros - write zero bits until byte boundary (0-7bits)
+func (w *EBSPWriter) StuffByteWithZeros() {
+	if w.n > 0 {
+		w.Write(0, 8-w.n)
+	}
+}

examples/initcreator/main.go

@@ -25,14 +25,15 @@ func main() {
 }
 
 func writeVideoAVCInitSegment() error {
-	spsNALU, _ := hex.DecodeString(sps1nalu)
+	sps, _ := hex.DecodeString(sps1nalu)
+	spsNALUs := [][]byte{sps}
 	pps, _ := hex.DecodeString(pps1nalu)
 	ppsNALUs := [][]byte{pps}
 
 	videoTimescale := uint32(180000)
 	init := mp4.CreateEmptyMP4Init(videoTimescale, "video", "und")
 	trak := init.Moov.Trak[0]
-	trak.SetAVCDescriptor("avc1", spsNALU, ppsNALUs)
+	trak.SetAVCDescriptor("avc1", spsNALUs, ppsNALUs)
 	width := trak.Mdia.Minf.Stbl.Stsd.AvcX.Width
 	height := trak.Mdia.Minf.Stbl.Stsd.AvcX.Height
 	if width != 1280 || height != 720 {