av/codec/h264/h264dec/cavlc.go

/*
DESCRIPTION
  cavlc.go provides utilities for context-adaptive variable-length coding
  for the parsing of H.264 syntax structure fields.

AUTHORS
  Saxon A. Nelson-Milton <saxon@ausocean.org>

LICENSE
  Copyright (C) 2019 the Australian Ocean Lab (AusOcean).

  It is free software: you can redistribute it and/or modify them
  under the terms of the GNU General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  It is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
  for more details.

  You should have received a copy of the GNU General Public License
  in gpl.txt.  If not, see http://www.gnu.org/licenses.
*/

package h264dec

import (
	"encoding/csv"
	"errors"
	"fmt"
	"os"
	"strconv"

	"bitbucket.org/ausocean/av/codec/h264/h264dec/bits"
)

// Initialize the CAVLC coeff_token mapping table.
func init() {
	const file = "coefftokenmap.csv"
	f, err := os.Open(file)
	if err != nil {
		panic(fmt.Sprintf("could not open coefftokenmap.csv file, failed with error: %v", err))
	}
	defer f.Close()

	lines, err := csv.NewReader(f).ReadAll()
	if err != nil {
		panic(fmt.Sprintf("could not read lines from coeftokenmap.csv file, failed with error: %v", err))
	}

	coeffTokenMap, err = formCoeffTokenMap(lines)
	if err != nil {
		panic(fmt.Sprintf("could not form coeff_token map, failed with err: %v", err))
	}
}

// The number of columns in the coeffTokenMap defined below. This is
// representative of the number of defined nC ranges defined in table 9-5.
const nColumns = 6

// coeffTokenMap will a representation of table 9-5 from the specifications, and
// is indexed as follows, coeffToken[ nC group ][ number of coeff_token leading
// zeros ][ value of coeff_token ][ 0 for TrailingOnes(coeff_token) and 1 for
// TotalCoef(coeff_token) ].
var coeffTokenMap [nColumns]map[int]map[int][2]int

// formCoeffTokenMap populates the global [nColumns]map[int]map[int][2]int,
// coeffTokenMap representation of table 9-5 from the specifications using lines
// read from a corresponding CSV file coefftokenmap.csv.
func formCoeffTokenMap(lines [][]string) ([nColumns]map[int]map[int][2]int, error) {
	var tokenMap [nColumns]map[int]map[int][2]int

	for i := range tokenMap {
		tokenMap[i] = make(map[int]map[int][2]int)
	}

	for _, line := range lines {
		trailingOnes, err := strconv.Atoi(line[0])
		if err != nil {
			return tokenMap, fmt.Errorf("could not convert trailingOnes string to int, failed with error: %v", err)
		}

		totalCoeff, err := strconv.Atoi(line[1])
		if err != nil {
			return tokenMap, fmt.Errorf("could not convert totalCoeff string to int, failed with error: %v", err)
		}

		// For each column in this row, therefore each nC category, load the
		// coeff_token leading zeros and value into the map.
		for j, v := range line[2:] {
			if v[0] == '-' {
				continue
			}

			// Count the leading zeros.
			var nZeros int
			for _, c := range v {
				if c == ' ' {
					continue
				}

				if c == '0' {
					nZeros++
					continue
				}
				break
			}

			// This will be the value of the coeff_token (without leading zeros).
			val, err := binToInt(v[nZeros:])
			if err != nil {
				return tokenMap, fmt.Errorf("could not get value of remaining binary, failed with error: %v", err)
			}

			// Add the TrailingOnes(coeff_token) and TotalCoeff(coeff_token) values
			// into the map for the coeff_token leading zeros and value.
			if tokenMap[j][nZeros] == nil {
				tokenMap[j][nZeros] = make(map[int][2]int)
			}
			tokenMap[j][nZeros][val] = [2]int{trailingOnes, totalCoeff}
		}
	}
	return tokenMap, nil
}

// parseLevelPrefix parses the level_prefix variable as specified by the process
// outlined in section 9.2.2.1 in the specifications.
func parseLevelPrefix(br *bits.BitReader) (int, error) {
	zeros := -1
	for b := 0; b != 1; zeros++ {
		_b, err := br.ReadBits(1)
		if err != nil {
			return -1, fmt.Errorf("could not read bit, failed with error: %v", err)
		}
		b = int(_b)
	}
	return zeros, nil
}

// parseLevelInformation parses level information and returns the resultant
// levelVal list using the process defined by section 9.2.2 in the specifications.
func parseLevelInformation(br *bits.BitReader, totalCoeff, trailingOnes int) ([]int, error) {
	var levelVal []int
	var i int
	for ; i < trailingOnes; i++ {
		b, err := br.ReadBits(1)
		if err != nil {
			return nil, fmt.Errorf("could not read trailing_ones_sign_flag, failed with error: %v", err)
		}
		levelVal = append(levelVal, 1-int(b)*2)
	}

	var suffixLen int
	switch {
	case totalCoeff > 10 && trailingOnes < 3:
		suffixLen = 1
	case totalCoeff <= 10 || trailingOnes == 3:
		suffixLen = 0
	default:
		return nil, errors.New("invalid TotalCoeff and TrailingOnes combination")
	}

	for j := 0; j < totalCoeff-trailingOnes; j++ {
		levelPrefix, err := parseLevelPrefix(br)
		if err != nil {
			return nil, fmt.Errorf("could not parse level prefix, failed with error: %v", err)
		}

		var levelSuffixSize int
		switch {
		case levelPrefix == 14 && suffixLen == 0:
			levelSuffixSize = 4
		case levelPrefix >= 15:
			levelSuffixSize = levelPrefix - 3
		default:
			levelSuffixSize = suffixLen
		}

		var levelSuffix int
		if levelSuffixSize > 0 {
			b, err := br.ReadBits(levelSuffixSize)
			if err != nil {
				return nil, fmt.Errorf("could not parse levelSuffix, failed with error: %v", err)
			}
			levelSuffix = int(b)
		} else {
			levelSuffix = 0
		}

		levelCode := (mini(15, levelPrefix) << uint(suffixLen)) + levelSuffix

		if levelPrefix >= 15 && suffixLen == 0 {
			levelCode += 15
		}

		if levelPrefix >= 16 {
			levelCode += (1 << uint(levelPrefix-3)) - 4096
		}

		if i == trailingOnes && trailingOnes < 3 {
			levelCode += 2
		}

		if levelCode%2 == 0 {
			levelVal = append(levelVal, (levelCode+2)>>1)
		} else {
			levelVal = append(levelVal, (-levelCode-1)>>1)
		}

		if suffixLen == 0 {
			suffixLen = 1
		}

		if absi(levelVal[i]) > (3<<uint(suffixLen-1)) && suffixLen < 6 {
			suffixLen++
		}
		i++
	}
	return levelVal, nil
}