diff --git a/avc/slice.go b/avc/slice.go index 4314702d..f4d73ae1 100644 --- a/avc/slice.go +++ b/avc/slice.go @@ -208,6 +208,7 @@ func ParseSliceHeader(nalu []byte, spsMap map[uint32]*SPS, ppsMap map[uint32]*PP if sliceType != SLICE_I && sliceType != SLICE_SI { sh.RefPicListModificationL0Flag = r.ReadFlag() if sh.RefPicListModificationL0Flag { + refPicListL0Loop: for { sh.ModificationOfPicNumsIDC = uint32(r.ReadExpGolomb()) switch sh.ModificationOfPicNumsIDC { @@ -216,10 +217,10 @@ func ParseSliceHeader(nalu []byte, spsMap map[uint32]*SPS, ppsMap map[uint32]*PP case 2: sh.LongTermPicNum = uint32(r.ReadExpGolomb()) case 3: - break + break refPicListL0Loop } if r.AccError() != nil { - break + break refPicListL0Loop } } } @@ -227,6 +228,7 @@ func ParseSliceHeader(nalu []byte, spsMap map[uint32]*SPS, ppsMap map[uint32]*PP if sliceType == SLICE_B { sh.RefPicListModificationL1Flag = r.ReadFlag() if sh.RefPicListModificationL1Flag { + refPicListL1Loop: for { sh.ModificationOfPicNumsIDC = uint32(r.ReadExpGolomb()) switch sh.ModificationOfPicNumsIDC { @@ -235,7 +237,10 @@ func ParseSliceHeader(nalu []byte, spsMap map[uint32]*SPS, ppsMap map[uint32]*PP case 2: sh.LongTermPicNum = uint32(r.ReadExpGolomb()) case 3: - break + break refPicListL1Loop + } + if r.AccError() != nil { + break refPicListL1Loop } } } @@ -299,6 +304,7 @@ func ParseSliceHeader(nalu []byte, spsMap map[uint32]*SPS, ppsMap map[uint32]*PP } else { sh.AdaptiveRefPicMarkingModeFlag = r.ReadFlag() if sh.AdaptiveRefPicMarkingModeFlag { + adaptiveRefPicLoop: for { memoryManagementControlOperation := r.ReadExpGolomb() switch memoryManagementControlOperation { @@ -313,7 +319,10 @@ func ParseSliceHeader(nalu []byte, spsMap map[uint32]*SPS, ppsMap map[uint32]*PP case 4: sh.MaxLongTermFrameIdxPlus1 = uint32(r.ReadExpGolomb()) case 0: - break + break adaptiveRefPicLoop + } + if r.AccError() != nil { + break adaptiveRefPicLoop } } } diff --git a/avc/slice_test.go b/avc/slice_test.go index 0890c8a2..55316098 100644 --- a/avc/slice_test.go +++ b/avc/slice_test.go @@ -23,7 +23,7 @@ func TestSliceTypeParser(t *testing.T) { } } -func TestParseSliceHeader(t *testing.T) { +func TestParseSliceHeader_BlackFrame(t *testing.T) { wantedHdr := SliceHeader{ SliceType: 7, SliceQPDelta: 6, @@ -65,3 +65,55 @@ func TestParseSliceHeader(t *testing.T) { t.Error(diff) } } + +func TestParseSliceHeader_TwoFrames(t *testing.T) { + wantedIdrHdr := SliceHeader{SliceType: SLICE_I, IDRPicID: 1, SliceQPDelta: 8, Size: 5} + wantedNonIdrHdr := SliceHeader{SliceType: SLICE_P, FrameNum: 1, ModificationOfPicNumsIDC: 3, SliceQPDelta: 13, Size: 5, NumRefIdxActiveOverrideFlag: true, RefPicListModificationL0Flag: true} + + data, err := ioutil.ReadFile("testdata/two-frames.264") + if err != nil { + t.Error(err) + } + nalus, err := GetNalusFromSample(data) + if err != nil { + t.Error(err) + } + spsMap := make(map[uint32]*SPS, 1) + ppsMap := make(map[uint32]*PPS, 1) + var gotIdrHdr *SliceHeader + var gotNonIdrHdr *SliceHeader + for _, nalu := range nalus { + switch GetNaluType(nalu[0]) { + case NALU_SPS: + sps, err := ParseSPSNALUnit(nalu, true) + if err != nil { + t.Error(err) + } + spsMap[uint32(sps.ParameterID)] = sps + case NALU_PPS: + pps, err := ParsePPSNALUnit(nalu, spsMap) + if err != nil { + t.Error(err) + } + ppsMap[uint32(pps.PicParameterSetID)] = pps + case NALU_IDR: + gotIdrHdr, err = ParseSliceHeader(nalu, spsMap, ppsMap) + if err != nil { + t.Error(err) + } + case NALU_NON_IDR: + gotNonIdrHdr, err = ParseSliceHeader(nalu, spsMap, ppsMap) + if err != nil { + t.Error(err) + } + } + } + if diff := deep.Equal(wantedIdrHdr, *gotIdrHdr); diff != nil { + fmt.Printf("Got IDR Slice Header: %+v\n Diff is: ", *gotIdrHdr) + t.Error(diff) + } + if diff := deep.Equal(wantedNonIdrHdr, *gotNonIdrHdr); diff != nil { + fmt.Printf("Got NON_IDR Slice Header: %+v\n Diff is: ", *gotNonIdrHdr) + t.Error(diff) + } +} diff --git a/avc/testdata/two-frames.264 b/avc/testdata/two-frames.264 new file mode 100644 index 00000000..6d7dec4a Binary files /dev/null and b/avc/testdata/two-frames.264 differ