codec/h264/h264dec: added TestNewPredWeightTable

codec/h264/h264dec/helpers.go

@@ -19,7 +19,7 @@ func binToSlice(s string) ([]byte, error) {
 		bytes []byte
 	)
 
-	for _, c := range s {
+	for i, c := range s {
 		switch c {
 		case ' ':
 			continue
@@ -31,7 +31,7 @@ func binToSlice(s string) ([]byte, error) {
 		}
 
 		a >>= 1
-		if a == 0 {
+		if a == 0 || i == (len(s)-1) {
 			bytes = append(bytes, cur)
 			cur = 0
 			a = 0x80

codec/h264/h264dec/slice.go

@@ -30,6 +30,8 @@ type VideoStream struct {
 	SPS    *SPS
 	PPS    *PPS
 	Slices []*SliceContext
+
+	ChromaArrayType int
 }
 type SliceContext struct {
 	*NALUnit
@@ -121,7 +123,6 @@ func NewRefPicListModification(br *bits.BitReader, p *PPS, s *SliceHeader) (*Ref
 type PredWeightTable struct {
 	LumaLog2WeightDenom   int
 	ChromaLog2WeightDenom int
-	ChromaArrayType       int
 	LumaWeightL0Flag      bool
 	LumaWeightL0          []int
 	LumaOffsetL0          []int
@@ -139,7 +140,7 @@ type PredWeightTable struct {
 // NewPredWeightTable parses elements of a pred_weight_table following the
 // syntax structure defined in section 7.3.3.2, and returns as a new
 // PredWeightTable.
-func NewPredWeightTable(br *bits.BitReader, h *SliceHeader) (*PredWeightTable, error) {
+func NewPredWeightTable(br *bits.BitReader, h *SliceHeader, chromaArrayType int) (*PredWeightTable, error) {
 	p := &PredWeightTable{}
 	var err error
 
@@ -148,7 +149,7 @@ func NewPredWeightTable(br *bits.BitReader, h *SliceHeader) (*PredWeightTable, e
 		return nil, errors.Wrap(err, "could not parse LumaLog2WeightDenom")
 	}
 
-	if p.ChromaArrayType != 0 {
+	if chromaArrayType != 0 {
 		p.ChromaLog2WeightDenom, err = readUe(br)
 		if err != nil {
 			return nil, errors.Wrap(err, "could not parse ChromaLog2WeightDenom")
@@ -174,7 +175,7 @@ func NewPredWeightTable(br *bits.BitReader, h *SliceHeader) (*PredWeightTable, e
 			}
 			p.LumaOffsetL0 = append(p.LumaOffsetL0, se)
 		}
-		if p.ChromaArrayType != 0 {
+		if chromaArrayType != 0 {
 			b, err := br.ReadBits(1)
 			if err != nil {
 				return nil, errors.Wrap(err, "could not read ChromaWeightL0Flag")
@@ -221,7 +222,7 @@ func NewPredWeightTable(br *bits.BitReader, h *SliceHeader) (*PredWeightTable, e
 				}
 				p.LumaOffsetL1 = append(p.LumaOffsetL1, se)
 			}
-			if p.ChromaArrayType != 0 {
+			if chromaArrayType != 0 {
 				b, err := br.ReadBits(1)
 				if err != nil {
 					return nil, errors.Wrap(err, "could not read ChromaWeightL1Flag")
@@ -529,7 +530,7 @@ func NumMbPart(nalUnit *NALUnit, sps *SPS, header *SliceHeader, data *SliceData)
 	return numMbPart
 }
 
-func MbPred(sliceContext *SliceContext, br *bits.BitReader, rbsp []byte) error {
+func MbPred(chromaArrayType int, sliceContext *SliceContext, br *bits.BitReader, rbsp []byte) error {
 	var cabac *CABAC
 	sliceType := sliceTypeMap[sliceContext.Slice.Header.SliceType]
 	mbPartPredMode, err := MbPartPredMode(sliceContext.Slice.Data, sliceType, sliceContext.Slice.Data.MbType, 0)
@@ -630,7 +631,7 @@ func MbPred(sliceContext *SliceContext, br *bits.BitReader, rbsp []byte) error {
 			}
 
 		}
-		if sliceContext.Slice.Header.ChromaArrayType == 1 || sliceContext.Slice.Header.ChromaArrayType == 2 {
+		if chromaArrayType == 1 || chromaArrayType == 2 {
 			if sliceContext.PPS.EntropyCodingMode == 1 {
 				// TODO: ue(v) or ae(v)
 				binarization := NewBinarization(
@@ -882,7 +883,7 @@ func MbaffFrameFlag(sps *SPS, header *SliceHeader) int {
 	return 0
 }
 
-func NewSliceData(sliceContext *SliceContext, br *bits.BitReader) (*SliceData, error) {
+func NewSliceData(chromaArrayType int, sliceContext *SliceContext, br *bits.BitReader) (*SliceData, error) {
 	var cabac *CABAC
 	var err error
 	sliceContext.Slice.Data = &SliceData{BitReader: br}
@@ -1117,7 +1118,7 @@ func NewSliceData(sliceContext *SliceContext, br *bits.BitReader) (*SliceData, e
 						}
 					}
 					// TODO: fix nil argument for.
-					MbPred(sliceContext, br, nil)
+					MbPred(chromaArrayType, sliceContext, br, nil)
 				}
 				m, err = MbPartPredMode(sliceContext.Slice.Data, sliceContext.Slice.Data.SliceTypeName, sliceContext.Slice.Data.MbType, 0)
 				if err != nil {
@@ -1136,7 +1137,7 @@ func NewSliceData(sliceContext *SliceContext, br *bits.BitReader) (*SliceData, e
 					} else {
 						me, _ := readMe(
 							br,
-							uint(sliceContext.Slice.Header.ChromaArrayType),
+							uint(chromaArrayType),
 							// TODO: fix this
 							//MbPartPredMode(sliceContext.Slice.Data, sliceContext.Slice.Data.SliceTypeName, sliceContext.Slice.Data.MbType, 0)))
 							0)
@@ -1213,10 +1214,10 @@ func NewSliceData(sliceContext *SliceContext, br *bits.BitReader) (*SliceData, e
 func (c *SliceContext) Update(header *SliceHeader, data *SliceData) {
 	c.Slice = &Slice{Header: header, Data: data}
 }
-func NewSliceContext(videoStream *VideoStream, nalUnit *NALUnit, rbsp []byte, showPacket bool) (*SliceContext, error) {
+func NewSliceContext(vid *VideoStream, nalUnit *NALUnit, rbsp []byte, showPacket bool) (*SliceContext, error) {
 	var err error
-	sps := videoStream.SPS
+	sps := vid.SPS
-	pps := videoStream.PPS
+	pps := vid.PPS
 	logger.Printf("debug: %s RBSP %d bytes %d bits == \n", NALUnitType[int(nalUnit.Type)], len(rbsp), len(rbsp)*8)
 	logger.Printf("debug: \t%#v\n", rbsp[0:8])
 	var idrPic bool
@@ -1225,9 +1226,9 @@ func NewSliceContext(videoStream *VideoStream, nalUnit *NALUnit, rbsp []byte, sh
 	}
 	header := SliceHeader{}
 	if sps.UseSeparateColorPlane {
-		header.ChromaArrayType = 0
+		vid.ChromaArrayType = 0
 	} else {
-		header.ChromaArrayType = sps.ChromaFormat
+		vid.ChromaArrayType = sps.ChromaFormat
 	}
 	br := bits.NewBitReader(bytes.NewReader(rbsp))
 
@@ -1350,7 +1351,7 @@ func NewSliceContext(videoStream *VideoStream, nalUnit *NALUnit, rbsp []byte, sh
 	}
 
 	if (pps.WeightedPred && (sliceType == "P" || sliceType == "SP")) || (pps.WeightedBipred == 1 && sliceType == "B") {
-		header.PredWeightTable, err = NewPredWeightTable(br, &header)
+		header.PredWeightTable, err = NewPredWeightTable(br, &header, vid.ChromaArrayType)
 		if err != nil {
 			return nil, errors.Wrap(err, "could not parse PredWeightTable")
 		}
@@ -1420,7 +1421,7 @@ func NewSliceContext(videoStream *VideoStream, nalUnit *NALUnit, rbsp []byte, sh
 			Header: &header,
 		},
 	}
-	sliceContext.Slice.Data, err = NewSliceData(sliceContext, br)
+	sliceContext.Slice.Data, err = NewSliceData(vid.ChromaArrayType,sliceContext, br)
 	if err != nil {
 		return nil, errors.Wrap(err, "could not create slice data")
 	}

codec/h264/h264dec/slice_test.go

@@ -85,10 +85,7 @@ func TestNewRefPicListModification(t *testing.T) {
 				"010" + // ue(v) long_term_pic_num = 1
 
 				// Fourth modification does not exist
-				"00100" + // ue(v) modification_of_pic_nums_idc[0][3] = 3
+				"00100", // ue(v) modification_of_pic_nums_idc[0][3] = 3
-
-				// Padding bits
-				"00",
 
 			s: SliceHeader{
 				SliceType: 3,
@@ -127,17 +124,115 @@ func TestNewRefPicListModification(t *testing.T) {
 func TestNewPredWeightTable(t *testing.T) {
 	tests := []struct {
 		in              string
-		s SliceHeader
+		sliceHeader     SliceHeader
+		chromaArrayType int // ChromaArrayType
 		want            PredWeightTable
 	}{
 		{
 			in: "011" + // ue(v) luma_log2_weight_denom = 2
 				"00100" + // ue(v) chroma_log2_weigght_denom = 3
 
+				// list0
+				// i = 0
 				"1" + // u(1) luma_weight_l0_flag = true
 				"011" + // se(v) luma_weight_l0[0] = -1
 				"010" + // se(v) luma_offset_l0[0] = 1
 				"1" + // u(1) chroma_weight_l0_flag = true
+
+				// i = 0, j = 0
+				"010" + // se(v) chroma_weight_l0[0][0] = 1
+				"010" + // se(v) chroma_offset_l0[0][0] = 1
+
+				// i = 0, j = 1
+				"010" + // se(v) chroma_weight_l0[0][1] = 1
+				"010" + // se(v) chroma_offset_l0[0][1] = 1
+
+				// i = 1
+				"1" + // u(1) luma_weight_l0_flag = true
+				"011" + // se(v) luma_weight_l0[1] = -1
+				"00100" + // se(v) luma_offset_l0[1] = 2
+				"1" + // u(1) chroma_weight_l0_flag = true
+
+				// i = 1, j = 0
+				"011" + // se(v) chroma_weight_l0[1][0] = -1
+				"00101" + // se(v) chroma_offset_l0[1][0] = -2
+
+				// i = 1, j = 1
+				"011" + // se(v) chroma_weight_l0[1][1] = -1
+				"011" + // se(v) chroma_offset_l0[1][1] = -1
+
+				// list1
+				// i = 0
+				"1" + // u(1) luma_weight_l1_flag = true
+				"011" + // se(v) luma_weight_l1[0] = -1
+				"010" + // se(v) luma_offset_l1[0] = 1
+				"1" + // u(1) chroma_weight_l1_flag = true
+
+				// i = 0, j = 0
+				"010" + // se(v) chroma_weight_l1[0][0] = 1
+				"010" + // se(v) chroma_offset_l1[0][0] = 1
+
+				// i = 0, j = 1
+				"010" + // se(v) chroma_weight_l1[0][1] = 1
+				"010" + // se(v) chroma_offset_l1[0][1] = 1
+
+				// i = 1
+				"1" + // u(1) luma_weight_l1_flag = true
+				"011" + // se(v) luma_weight_l1[1] = -1
+				"00100" + // se(v) luma_offset_l1[1] = 2
+				"1" + // u(1) chroma_weight_l1_flag = true
+
+				// i = 1, j = 0
+				"011" + // se(v) chroma_weight_l0[1][0] = -1
+				"00101" + // se(v) chroma_offset_l0[1][0] = -2
+
+				// i = 1, j = 1
+				"011" + // se(v) chroma_weight_l0[1][1] = -1
+				"011", // se(v) chroma_offset_l0[1][1] = -1
+
+			sliceHeader: SliceHeader{
+				NumRefIdxL0ActiveMinus1: 1,
+				NumRefIdxL1ActiveMinus1: 1,
+				SliceType:               1,
+			},
+
+			chromaArrayType: 1,
+
+			want: PredWeightTable{
+				LumaLog2WeightDenom:   2,
+				ChromaLog2WeightDenom: 3,
+
+				LumaWeightL0Flag:   true,
+				LumaWeightL0:       []int{-1, -1},
+				LumaOffsetL0:       []int{1, 2},
+				ChromaWeightL0Flag: true,
+				ChromaWeightL0:     [][]int{{1, 1}, {-1, -1}},
+				ChromaOffsetL0:     [][]int{{1, 1}, {-2, -1}},
+
+				LumaWeightL1Flag:   true,
+				LumaWeightL1:       []int{-1, -1},
+				LumaOffsetL1:       []int{1, 2},
+				ChromaWeightL1Flag: true,
+				ChromaWeightL1:     [][]int{{1, 1}, {-1, -1}},
+				ChromaOffsetL1:     [][]int{{1, 1}, {-2, -1}},
+			},
 		},
 	}
+
+	for i, test := range tests {
+		inBytes, err := binToSlice(test.in)
+		if err != nil {
+			t.Fatalf("unexpected error %v for binToSlice in test %d", err, i)
+		}
+
+		got, err := NewPredWeightTable(bits.NewBitReader(bytes.NewReader(inBytes)),
+			&test.sliceHeader, test.chromaArrayType)
+		if err != nil {
+			t.Fatalf("unexpected error %v for NewPredWeightTable in test %d", err, i)
+		}
+
+		if !reflect.DeepEqual(*got, test.want) {
+			t.Errorf("did not get expected result for test %d\nGot: %v\nWant: %v\n", i, got, test.want)
+		}
+	}
 }