1 /*
2  * Copyright (c) 2023 Huawei Device Co., Ltd.
3  * Licensed under the Apache License, Version 2.0 (the "License");
4  * you may not use this file except in compliance with the License.
5  * You may obtain a copy of the License at
6  *
7  *     http://www.apache.org/licenses/LICENSE-2.0
8  *
9  * Unless required by applicable law or agreed to in writing, software
10  * distributed under the License is distributed on an "AS IS" BASIS,
11  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12  * See the License for the specific language governing permissions and
13  * limitations under the License.
14  */
15 
16 #ifndef UTIL_STARTCODEDETECTOR_H
17 #define UTIL_STARTCODEDETECTOR_H
18 
19 #include <string>
20 #include <list>
21 #include <vector>
22 #include <optional>
23 #include <fstream>
24 #include <memory>
25 
26 enum CodeType {
27     H264,
28     H265,
29     H266,
30 };
31 
32 struct Sample {
33     size_t startPos;
34     size_t endPos;
35     bool isCsd;
36     bool isIdr;
37     size_t idx;
38     std::string s;
39     std::vector<uint8_t> vividSei;
40 };
41 
42 class StartCodeDetector {
43 public:
44     static std::shared_ptr<StartCodeDetector> Create(CodeType type);
45     size_t SetSource(const std::string &path);  // return sample cnt
46     size_t SetSource(const uint8_t* pStart, size_t bufSize);  // return sample cnt
47     bool SeekTo(size_t sampleIdx);
48     std::optional<Sample> PeekNextSample();
49     void MoveToNext();
50 
51 protected:
52     StartCodeDetector() = default;
53     virtual ~StartCodeDetector() = default;
54 
55 private:
56     struct NALUInfo {
57         size_t startPos;
58         size_t endPos;
59         uint8_t nalType;
60         std::vector<uint8_t> vividSei;
61     };
62     void SaveVivid(NALUInfo& nal, const uint8_t *pStart);
63     void BuildSampleList();
64     static size_t GetFileSizeInBytes(std::ifstream &ifs);
65 
66     virtual uint8_t GetNalType(uint8_t firstByte, uint8_t secondByte) = 0;
67     virtual bool IsPPS(uint8_t nalType) = 0;
68     virtual bool IsVCL(uint8_t nalType) = 0;
69     virtual bool IsIDR(uint8_t nalType) = 0;
IsPrefixSEI(uint8_t nalType)70     virtual bool IsPrefixSEI(uint8_t nalType) { return false; }
71 
72     static constexpr uint8_t START_CODE[] = {0, 0, 1};
73     static constexpr size_t START_CODE_LEN = sizeof(START_CODE);
74 
75     std::list<NALUInfo> nals_;
76     std::vector<Sample> samples_;
77     std::list<size_t> csdIdxList_;
78     std::list<size_t> idrIdxList_;
79     std::optional<size_t> waitingCsd_;  // if user seek previously, the csd sample will be stored here
80     size_t nextSampleIdx_ = 0;
81 };
82 
83 class StartCodeDetectorH264 : public StartCodeDetector {
84 public:
85     StartCodeDetectorH264() = default;
86     ~StartCodeDetectorH264() override = default;
87 
88 private:
89     enum H264NalType : uint8_t {
90         UNSPECIFIED = 0,
91         NON_IDR = 1,
92         PARTITION_A = 2,
93         PARTITION_B = 3,
94         PARTITION_C = 4,
95         IDR = 5,
96         SEI = 6,
97         SPS = 7,
98         PPS = 8,
99         AU_DELIMITER = 9,
100         END_OF_SEQUENCE = 10,
101         END_OF_STREAM = 11,
102         FILLER_DATA = 12,
103         SPS_EXT = 13,
104         PREFIX = 14,
105         SUB_SPS = 15,
106         DPS = 16,
107     };
108     uint8_t GetNalType(uint8_t firstByte, uint8_t) override;
109     bool IsPPS(uint8_t nalType) override;
110     bool IsVCL(uint8_t nalType) override;
111     bool IsIDR(uint8_t nalType) override;
112 };
113 
114 class StartCodeDetectorH265 : public StartCodeDetector {
115 public:
116     StartCodeDetectorH265() = default;
117     ~StartCodeDetectorH265() override = default;
118 
119 private:
120     enum H265NalType : uint8_t {
121         HEVC_TRAIL_N = 0,
122         HEVC_TRAIL_R = 1,
123         HEVC_TSA_N = 2,
124         HEVC_TSA_R = 3,
125         HEVC_STSA_N = 4,
126         HEVC_STSA_R = 5,
127         HEVC_RADL_N = 6,
128         HEVC_RADL_R = 7,
129         HEVC_RASL_N = 8,
130         HEVC_RASL_R = 9,
131         HEVC_BLA_W_LP = 16,
132         HEVC_BLA_W_RADL = 17,
133         HEVC_BLA_N_LP = 18,
134         HEVC_IDR_W_RADL = 19,
135         HEVC_IDR_N_LP = 20,
136         HEVC_CRA_NUT = 21,
137         HEVC_VPS_NUT = 32,
138         HEVC_SPS_NUT = 33,
139         HEVC_PPS_NUT = 34,
140         HEVC_AUD_NUT = 35,
141         HEVC_EOS_NUT = 36,
142         HEVC_EOB_NUT = 37,
143         HEVC_FD_NUT = 38,
144         HEVC_PREFIX_SEI_NUT = 39,
145         HEVC_SUFFIX_SEI_NUT = 40,
146     };
147     uint8_t GetNalType(uint8_t firstByte, uint8_t) override;
148     bool IsPPS(uint8_t nalType) override;
149     bool IsVCL(uint8_t nalType) override;
150     bool IsIDR(uint8_t nalType) override;
151     bool IsPrefixSEI(uint8_t nalType) override;
152 };
153 
154 class StartCodeDetectorH266 : public StartCodeDetector {
155 public:
156     StartCodeDetectorH266() = default;
157     ~StartCodeDetectorH266() override = default;
158 
159 private:
160     enum H266NalType : uint8_t {
161         // VCL
162         VVC_TRAIL_NUT = 0,
163         VVC_STSA_NUT = 1,
164         VVC_RADL_NUT = 2,
165         VVC_RASL_NUT = 3,
166         VVC_RSV_VCL_4 = 4,
167         VVC_RSV_VCL_5 = 5,
168         VVC_RSV_VCL_6 = 6,
169         VVC_IDR_W_RADL = 7,
170         VVC_IDR_N_LP = 8,
171         VVC_CRA_NUT = 9,
172         VVC_GDR_NUT = 10,
173         VVC_RSV_IRAP_11 = 11,
174         // non-VCL
175         VVC_OPI_NUT = 12,
176         VVC_DCI_NUT = 13,
177         VVC_VPS_NUT = 14,
178         VVC_SPS_NUT = 15,
179         VVC_PPS_NUT = 16,
180         VVC_PREFIX_APS_NUT = 17,
181         VVC_SUFFIX_APS_NUT = 18,
182         VVC_PH_NUT = 19,
183         VVC_AUD_NUT = 20,
184         VVC_EOS_NUT = 21,
185         VVC_EOB_NUT = 22,
186         VVC_PREFIX_SEI_NUT = 23,
187         VVC_SUFFIX_SEI_NUT = 24,
188         VVC_FD_NUT = 25,
189     };
190     uint8_t GetNalType(uint8_t, uint8_t secondByte) override;
191     bool IsPPS(uint8_t nalType) override;
192     bool IsVCL(uint8_t nalType) override;
193     bool IsIDR(uint8_t nalType) override;
194     bool IsPrefixSEI(uint8_t nalType) override;
195 };
196 #endif // UTIL_STARTCODEDETECTOR_H
197