github.com/bluenviron/mediacommon@v1.9.3/pkg/codecs/h264/annexb.go (about)

     1  package h264
     2  
     3  import (
     4  	"fmt"
     5  )
     6  
     7  // AnnexBUnmarshal decodes an access unit from the Annex-B stream format.
     8  // Specification: ITU-T Rec. H.264, Annex B
     9  func AnnexBUnmarshal(buf []byte) ([][]byte, error) {
    10  	bl := len(buf)
    11  	initZeroCount := 0
    12  	start := 0
    13  
    14  outer:
    15  	for {
    16  		if start >= bl || start >= 4 {
    17  			return nil, fmt.Errorf("initial delimiter not found")
    18  		}
    19  
    20  		switch initZeroCount {
    21  		case 0, 1:
    22  			if buf[start] != 0 {
    23  				return nil, fmt.Errorf("initial delimiter not found")
    24  			}
    25  			initZeroCount++
    26  
    27  		case 2, 3:
    28  			switch buf[start] {
    29  			case 1:
    30  				start++
    31  				break outer
    32  
    33  			case 0:
    34  
    35  			default:
    36  				return nil, fmt.Errorf("initial delimiter not found")
    37  			}
    38  			initZeroCount++
    39  		}
    40  
    41  		start++
    42  	}
    43  
    44  	zeroCount := 0
    45  	n := 0
    46  
    47  	for i := start; i < bl; i++ {
    48  		switch buf[i] {
    49  		case 0:
    50  			zeroCount++
    51  
    52  		case 1:
    53  			if zeroCount == 2 || zeroCount == 3 {
    54  				n++
    55  			}
    56  			zeroCount = 0
    57  
    58  		default:
    59  			zeroCount = 0
    60  		}
    61  	}
    62  
    63  	if (n + 1) > MaxNALUsPerAccessUnit {
    64  		return nil, fmt.Errorf("NALU count (%d) exceeds maximum allowed (%d)",
    65  			n+1, MaxNALUsPerAccessUnit)
    66  	}
    67  
    68  	ret := make([][]byte, n+1)
    69  	pos := 0
    70  	start = initZeroCount + 1
    71  	zeroCount = 0
    72  	delimStart := 0
    73  	auSize := 0
    74  
    75  	for i := start; i < bl; i++ {
    76  		switch buf[i] {
    77  		case 0:
    78  			if zeroCount == 0 {
    79  				delimStart = i
    80  			}
    81  			zeroCount++
    82  
    83  		case 1:
    84  			if zeroCount == 2 || zeroCount == 3 {
    85  				l := delimStart - start
    86  
    87  				if l == 0 {
    88  					return nil, fmt.Errorf("invalid NALU")
    89  				}
    90  
    91  				if (auSize + l) > MaxAccessUnitSize {
    92  					return nil, fmt.Errorf("access unit size (%d) is too big, maximum is %d", auSize+l, MaxAccessUnitSize)
    93  				}
    94  
    95  				ret[pos] = buf[start:delimStart]
    96  				pos++
    97  				auSize += l
    98  				start = i + 1
    99  			}
   100  			zeroCount = 0
   101  
   102  		default:
   103  			zeroCount = 0
   104  		}
   105  	}
   106  
   107  	l := bl - start
   108  
   109  	if l == 0 {
   110  		return nil, fmt.Errorf("invalid NALU")
   111  	}
   112  
   113  	if (auSize + l) > MaxAccessUnitSize {
   114  		return nil, fmt.Errorf("access unit size (%d) is too big, maximum is %d", auSize+l, MaxAccessUnitSize)
   115  	}
   116  
   117  	ret[pos] = buf[start:bl]
   118  
   119  	return ret, nil
   120  }
   121  
   122  func annexBMarshalSize(au [][]byte) int {
   123  	n := 0
   124  	for _, nalu := range au {
   125  		n += 4 + len(nalu)
   126  	}
   127  	return n
   128  }
   129  
   130  // AnnexBMarshal encodes an access unit into the Annex-B stream format.
   131  // Specification: ITU-T Rec. H.264, Annex B
   132  func AnnexBMarshal(au [][]byte) ([]byte, error) {
   133  	buf := make([]byte, annexBMarshalSize(au))
   134  	pos := 0
   135  
   136  	for _, nalu := range au {
   137  		pos += copy(buf[pos:], []byte{0x00, 0x00, 0x00, 0x01})
   138  		pos += copy(buf[pos:], nalu)
   139  	}
   140  
   141  	return buf, nil
   142  }