|
| 1 | +// Package h265reader implements a H265 Annex-B Reader |
| 2 | +package gb28181 |
| 3 | + |
| 4 | +import ( |
| 5 | + "bytes" |
| 6 | + "errors" |
| 7 | + "io" |
| 8 | +) |
| 9 | + |
| 10 | +type NalUnitType uint8 |
| 11 | + |
| 12 | +// Enums for NalUnitTypes |
| 13 | +const ( |
| 14 | + NaluTypeSliceTrailN NalUnitType = 0 // 0x0 |
| 15 | + NaluTypeSliceTrailR NalUnitType = 1 // 0x01 |
| 16 | + NaluTypeSliceTsaN NalUnitType = 2 // 0x02 |
| 17 | + NaluTypeSliceTsaR NalUnitType = 3 // 0x03 |
| 18 | + NaluTypeSliceStsaN NalUnitType = 4 // 0x04 |
| 19 | + NaluTypeSliceStsaR NalUnitType = 5 // 0x05 |
| 20 | + NaluTypeSliceRadlN NalUnitType = 6 // 0x06 |
| 21 | + NaluTypeSliceRadlR NalUnitType = 7 // 0x07 |
| 22 | + NaluTypeSliceRaslN NalUnitType = 8 // 0x06 |
| 23 | + NaluTypeSliceRaslR NalUnitType = 9 // 0x09 |
| 24 | + |
| 25 | + NaluTypeSliceBlaWlp NalUnitType = 16 // 0x10 |
| 26 | + NaluTypeSliceBlaWradl NalUnitType = 17 // 0x11 |
| 27 | + NaluTypeSliceBlaNlp NalUnitType = 18 // 0x12 |
| 28 | + NaluTypeSliceIdr NalUnitType = 19 // 0x13 |
| 29 | + NaluTypeSliceIdrNlp NalUnitType = 20 // 0x14 |
| 30 | + NaluTypeSliceCranut NalUnitType = 21 // 0x15 |
| 31 | + NaluTypeSliceRsvIrapVcl22 NalUnitType = 22 // 0x16 |
| 32 | + NaluTypeSliceRsvIrapVcl23 NalUnitType = 23 // 0x17 |
| 33 | + |
| 34 | + NaluTypeVps NalUnitType = 32 // 0x20 |
| 35 | + NaluTypeSps NalUnitType = 33 // 0x21 |
| 36 | + NaluTypePps NalUnitType = 34 // 0x22 |
| 37 | + NaluTypeAud NalUnitType = 35 // 0x23 |
| 38 | + NaluTypeSei NalUnitType = 39 // 0x27 |
| 39 | + NaluTypeSeiSuffix NalUnitType = 40 // 0x28 |
| 40 | + |
| 41 | + NaluTypeUnspecified NalUnitType = 48 // 0x30 |
| 42 | +) |
| 43 | + |
| 44 | +// H265Reader reads data from stream and constructs h265 nal units |
| 45 | +type H265Reader struct { |
| 46 | + stream io.Reader |
| 47 | + nalBuffer []byte |
| 48 | + countOfConsecutiveZeroBytes int |
| 49 | + nalPrefixParsed bool |
| 50 | + readBuffer []byte |
| 51 | +} |
| 52 | + |
| 53 | +var ( |
| 54 | + errNilReader = errors.New("stream is nil") |
| 55 | + errDataIsNotH265Stream = errors.New("data is not a H265 bitstream") |
| 56 | +) |
| 57 | + |
| 58 | +// NewReader creates new H265Reader |
| 59 | +func NewReader(in io.Reader) (*H265Reader, error) { |
| 60 | + if in == nil { |
| 61 | + return nil, errNilReader |
| 62 | + } |
| 63 | + |
| 64 | + reader := &H265Reader{ |
| 65 | + stream: in, |
| 66 | + nalBuffer: make([]byte, 0), |
| 67 | + nalPrefixParsed: false, |
| 68 | + readBuffer: make([]byte, 0), |
| 69 | + } |
| 70 | + |
| 71 | + return reader, nil |
| 72 | +} |
| 73 | + |
| 74 | +// NAL H.265 Network Abstraction Layer |
| 75 | +type NAL struct { |
| 76 | + PictureOrderCount uint32 |
| 77 | + |
| 78 | + // NAL header |
| 79 | + ForbiddenZeroBit bool |
| 80 | + UnitType NalUnitType |
| 81 | + NuhLayerId uint8 |
| 82 | + NuhTemporalIdPlus1 uint8 |
| 83 | + |
| 84 | + Data []byte // header byte + rbsp |
| 85 | +} |
| 86 | + |
| 87 | +func (reader *H265Reader) read(numToRead int) (data []byte) { |
| 88 | + for len(reader.readBuffer) < numToRead { |
| 89 | + buf := make([]byte, 4096) |
| 90 | + n, err := reader.stream.Read(buf) |
| 91 | + if n == 0 || err != nil { |
| 92 | + break |
| 93 | + } |
| 94 | + buf = buf[0:n] |
| 95 | + reader.readBuffer = append(reader.readBuffer, buf...) |
| 96 | + } |
| 97 | + var numShouldRead int |
| 98 | + if numToRead <= len(reader.readBuffer) { |
| 99 | + numShouldRead = numToRead |
| 100 | + } else { |
| 101 | + numShouldRead = len(reader.readBuffer) |
| 102 | + } |
| 103 | + data = reader.readBuffer[0:numShouldRead] |
| 104 | + reader.readBuffer = reader.readBuffer[numShouldRead:] |
| 105 | + return data |
| 106 | +} |
| 107 | + |
| 108 | +func (reader *H265Reader) bitStreamStartsWithH265Prefix() (prefixLength int, e error) { |
| 109 | + nalPrefix3Bytes := []byte{0, 0, 1} |
| 110 | + nalPrefix4Bytes := []byte{0, 0, 0, 1} |
| 111 | + |
| 112 | + prefixBuffer := reader.read(4) |
| 113 | + |
| 114 | + n := len(prefixBuffer) |
| 115 | + |
| 116 | + if n == 0 { |
| 117 | + return 0, io.EOF |
| 118 | + } |
| 119 | + |
| 120 | + if n < 3 { |
| 121 | + return 0, errDataIsNotH265Stream |
| 122 | + } |
| 123 | + |
| 124 | + nalPrefix3BytesFound := bytes.Equal(nalPrefix3Bytes, prefixBuffer[:3]) |
| 125 | + if n == 3 { |
| 126 | + if nalPrefix3BytesFound { |
| 127 | + return 0, io.EOF |
| 128 | + } |
| 129 | + return 0, errDataIsNotH265Stream |
| 130 | + } |
| 131 | + |
| 132 | + // n == 4 |
| 133 | + if nalPrefix3BytesFound { |
| 134 | + reader.nalBuffer = append(reader.nalBuffer, prefixBuffer[3]) |
| 135 | + return 3, nil |
| 136 | + } |
| 137 | + |
| 138 | + nalPrefix4BytesFound := bytes.Equal(nalPrefix4Bytes, prefixBuffer) |
| 139 | + if nalPrefix4BytesFound { |
| 140 | + return 4, nil |
| 141 | + } |
| 142 | + return 0, errDataIsNotH265Stream |
| 143 | +} |
| 144 | + |
| 145 | +// NextNAL reads from stream and returns then next NAL, |
| 146 | +// and an error if there is incomplete frame data. |
| 147 | +// Returns all nil values when no more NALs are available. |
| 148 | +func (reader *H265Reader) NextNAL() (*NAL, error) { |
| 149 | + if !reader.nalPrefixParsed { |
| 150 | + _, err := reader.bitStreamStartsWithH265Prefix() |
| 151 | + if err != nil { |
| 152 | + return nil, err |
| 153 | + } |
| 154 | + |
| 155 | + reader.nalPrefixParsed = true |
| 156 | + } |
| 157 | + |
| 158 | + for { |
| 159 | + buffer := reader.read(1) |
| 160 | + n := len(buffer) |
| 161 | + |
| 162 | + if n != 1 { |
| 163 | + break |
| 164 | + } |
| 165 | + readByte := buffer[0] |
| 166 | + nalFound := reader.processByte(readByte) |
| 167 | + if nalFound { |
| 168 | + nal := newNal(reader.nalBuffer) |
| 169 | + nal.parseHeader() |
| 170 | + if nal.UnitType == NaluTypeSeiSuffix || nal.UnitType == NaluTypeSei { |
| 171 | + reader.nalBuffer = nil |
| 172 | + continue |
| 173 | + } else { |
| 174 | + break |
| 175 | + } |
| 176 | + } |
| 177 | + |
| 178 | + reader.nalBuffer = append(reader.nalBuffer, readByte) |
| 179 | + } |
| 180 | + |
| 181 | + if len(reader.nalBuffer) == 0 { |
| 182 | + return nil, io.EOF |
| 183 | + } |
| 184 | + |
| 185 | + nal := newNal(reader.nalBuffer) |
| 186 | + reader.nalBuffer = nil |
| 187 | + nal.parseHeader() |
| 188 | + |
| 189 | + return nal, nil |
| 190 | +} |
| 191 | + |
| 192 | +func (reader *H265Reader) processByte(readByte byte) (nalFound bool) { |
| 193 | + nalFound = false |
| 194 | + |
| 195 | + switch readByte { |
| 196 | + case 0: |
| 197 | + reader.countOfConsecutiveZeroBytes++ |
| 198 | + case 1: |
| 199 | + if reader.countOfConsecutiveZeroBytes >= 2 { |
| 200 | + countOfConsecutiveZeroBytesInPrefix := 2 |
| 201 | + if reader.countOfConsecutiveZeroBytes > 2 { |
| 202 | + countOfConsecutiveZeroBytesInPrefix = 3 |
| 203 | + } |
| 204 | + nalUnitLength := len(reader.nalBuffer) - countOfConsecutiveZeroBytesInPrefix |
| 205 | + reader.nalBuffer = reader.nalBuffer[0:nalUnitLength] |
| 206 | + reader.countOfConsecutiveZeroBytes = 0 |
| 207 | + nalFound = true |
| 208 | + } else { |
| 209 | + reader.countOfConsecutiveZeroBytes = 0 |
| 210 | + } |
| 211 | + default: |
| 212 | + reader.countOfConsecutiveZeroBytes = 0 |
| 213 | + } |
| 214 | + |
| 215 | + return nalFound |
| 216 | +} |
| 217 | + |
| 218 | +func newNal(data []byte) *NAL { |
| 219 | + return &NAL{PictureOrderCount: 0, ForbiddenZeroBit: false, UnitType: NaluTypeUnspecified, Data: data} |
| 220 | +} |
| 221 | + |
| 222 | +func (h *NAL) parseHeader() { |
| 223 | + firstByte := h.Data[0] |
| 224 | + h.ForbiddenZeroBit = (((firstByte & 0x80) >> 7) == 1) // 0x80 = 0b10000000 |
| 225 | + h.UnitType = NalUnitType((firstByte & 0x7E) >> 1) // 0x1F = 0b01111110 |
| 226 | +} |
0 commit comments