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 #include "start_code_detector.h"
17 #include <memory>
18 #include <algorithm>
19 #include "hcodec_log.h"
20 using namespace std;
21 
Create(CodeType type)22 std::shared_ptr<StartCodeDetector> StartCodeDetector::Create(CodeType type)
23 {
24     switch (type) {
25         case H264:
26             return make_shared<StartCodeDetectorH264>();
27         case H265:
28             return make_shared<StartCodeDetectorH265>();
29         case H266:
30             return make_shared<StartCodeDetectorH266>();
31         default:
32             return nullptr;
33     }
34 }
35 
SetSource(const std::string & path)36 size_t StartCodeDetector::SetSource(const std::string &path)
37 {
38     ifstream ifs(path, ios::binary);
39     if (!ifs.is_open()) {
40         TLOGE("cannot open %s", path.c_str());
41         return 0;
42     }
43     size_t fileSize = GetFileSizeInBytes(ifs);
44     unique_ptr<uint8_t[]> buf = make_unique<uint8_t[]>(fileSize);
45     ifs.read(reinterpret_cast<char *>(buf.get()), static_cast<std::streamsize>(fileSize));
46     return SetSource(buf.get(), fileSize);
47 }
48 
SetSource(const uint8_t * pStart,size_t bufSize)49 size_t StartCodeDetector::SetSource(const uint8_t* pStart, size_t bufSize)
50 {
51     if (pStart == nullptr) {
52         return 0;
53     }
54     list<tuple<size_t, uint8_t, uint8_t>> posOfFile;
55     size_t pos = 0;
56     while (pos < bufSize) {
57         auto pFound = search(pStart + pos, pStart + bufSize, begin(START_CODE), end(START_CODE));
58         pos = distance(pStart, pFound);
59         if (pos == bufSize || pos + START_CODE_LEN >= bufSize) { // 没找到或找到的起始码正好在文件末尾
60             break;
61         }
62         posOfFile.emplace_back(pos, pStart[pos + START_CODE_LEN], pStart[pos + START_CODE_LEN + 1]);
63         pos += START_CODE_LEN;
64     }
65     for (auto it = posOfFile.begin(); it != posOfFile.end(); ++it) {
66         auto nex = next(it);
67         NALUInfo nal {
68             .startPos = get<0>(*it),
69             .endPos = (nex == posOfFile.end()) ? (bufSize) : (get<0>(*nex)),
70             .nalType = GetNalType(get<1>(*it), get<2>(*it)),
71         };
72         SaveVivid(nal, pStart);
73         nals_.push_back(nal);
74     }
75     BuildSampleList();
76     return samples_.size();
77 }
78 
SaveVivid(NALUInfo & nal,const uint8_t * pStart)79 void StartCodeDetector::SaveVivid(NALUInfo& nal, const uint8_t *pStart)
80 {
81     if (!IsPrefixSEI(nal.nalType)) {
82         return;
83     }
84     const uint8_t *nalStart = pStart + nal.startPos;
85     if (*(nalStart + 5) == 0x04 &&  // 5: offset of last_payload_type_byte, 0x04: itu_t_t35
86         *(nalStart + 7) == 0x26 &&  // 7: offset of itu_t_t35_country_code, 0x26: value in spec
87         *(nalStart + 8) == 0x00 &&  // 8: offset of terminal_provide_code, 0x00: value in spec
88         *(nalStart + 9) == 0x04 &&  // 9: offset of terminal_provide_code, 0x04: value in spec
89         *(nalStart + 10) == 0x00 && // 10: offset of terminal_provide_oriented_code, 0x00: value in spec
90         *(nalStart + 11) == 0x05) { // 11: offset of terminal_provide_oriented_code, 0x05: value in spec
91         nal.vividSei = vector<uint8_t>(pStart + nal.startPos, pStart + nal.endPos);
92     }
93 }
94 
BuildSampleList()95 void StartCodeDetector::BuildSampleList()
96 {
97     shared_ptr<Sample> sample;
98     for (auto& nal : nals_) {
99         if (sample == nullptr) {
100             sample = make_shared<Sample>();
101             sample->startPos = nal.startPos;
102             sample->isCsd = false;
103             sample->isIdr = false;
104         }
105         sample->endPos = nal.endPos;
106         if (!sample->s.empty()) {
107             sample->s += "+";
108         }
109         sample->s += to_string(nal.nalType);
110         if (!nal.vividSei.empty()) {
111             sample->vividSei = std::move(nal.vividSei);
112         }
113 
114         bool isPPS = IsPPS(nal.nalType);
115         bool isVCL = IsVCL(nal.nalType);
116         bool isIDR = IsIDR(nal.nalType);
117         if (isPPS || isVCL) {  // should cut here and build one sample
118             if (isPPS) {
119                 sample->isCsd = true;
120                 csdIdxList_.push_back(samples_.size());
121             }
122             if (isIDR) {
123                 sample->isIdr = true;
124                 idrIdxList_.push_back(samples_.size());
125             }
126             sample->idx = samples_.size();
127             samples_.push_back(*sample);
128             sample.reset();
129         }
130     }
131 }
132 
GetFileSizeInBytes(ifstream & ifs)133 size_t StartCodeDetector::GetFileSizeInBytes(ifstream &ifs)
134 {
135     ifs.seekg(0, ifstream::end);
136     auto len = ifs.tellg();
137     ifs.seekg(0, ifstream::beg);
138     return static_cast<size_t>(len);
139 }
140 
SeekTo(size_t sampleIdx)141 bool StartCodeDetector::SeekTo(size_t sampleIdx)
142 {
143     if (sampleIdx >= samples_.size()) {
144         return false;
145     }
146 
147     auto idrIter = find_if(idrIdxList_.rbegin(), idrIdxList_.rend(), [sampleIdx](size_t idrIdx) {
148         return idrIdx <= sampleIdx;
149     });
150     if (idrIter == idrIdxList_.rend()) {
151         return false;
152     }
153     size_t idrIdx = *idrIter;
154 
155     auto csdIter = find_if(csdIdxList_.rbegin(), csdIdxList_.rend(), [idrIdx](size_t csdIdx) {
156         return csdIdx < idrIdx;
157     });
158     if (csdIter == csdIdxList_.rend()) {
159         return false;
160     }
161     size_t csdIdx = *csdIter;
162     waitingCsd_ = csdIdx;
163     nextSampleIdx_ = idrIdx;
164     TLOGI("csd idx=%zu, idr idx=%zu, target sample idx=%zu", csdIdx, idrIdx, sampleIdx);
165     return true;
166 }
167 
PeekNextSample()168 std::optional<Sample> StartCodeDetector::PeekNextSample()
169 {
170     if (waitingCsd_.has_value()) {
171         return samples_[waitingCsd_.value()];
172     }
173     if (nextSampleIdx_ >= samples_.size()) {
174         return std::nullopt;
175     }
176     return samples_[nextSampleIdx_];
177 }
178 
MoveToNext()179 void StartCodeDetector::MoveToNext()
180 {
181     if (waitingCsd_.has_value()) {
182         waitingCsd_ = nullopt;
183         return;
184     }
185     nextSampleIdx_++;
186 }
187 
GetNalType(uint8_t firstByte,uint8_t)188 uint8_t StartCodeDetectorH264::GetNalType(uint8_t firstByte, uint8_t)
189 {
190     return firstByte & 0b0001'1111;
191 }
192 
IsPPS(uint8_t nalType)193 bool StartCodeDetectorH264::IsPPS(uint8_t nalType)
194 {
195     return nalType == H264NalType::PPS;
196 }
197 
IsVCL(uint8_t nalType)198 bool StartCodeDetectorH264::IsVCL(uint8_t nalType)
199 {
200     return nalType >= H264NalType::NON_IDR && nalType <= H264NalType::IDR;
201 }
202 
IsIDR(uint8_t nalType)203 bool StartCodeDetectorH264::IsIDR(uint8_t nalType)
204 {
205     return nalType == H264NalType::IDR;
206 }
207 
GetNalType(uint8_t firstByte,uint8_t)208 uint8_t StartCodeDetectorH265::GetNalType(uint8_t firstByte, uint8_t)
209 {
210     return (firstByte & 0b0111'1110) >> 1;
211 }
212 
IsPPS(uint8_t nalType)213 bool StartCodeDetectorH265::IsPPS(uint8_t nalType)
214 {
215     return nalType == H265NalType::HEVC_PPS_NUT;
216 }
217 
IsVCL(uint8_t nalType)218 bool StartCodeDetectorH265::IsVCL(uint8_t nalType)
219 {
220     return nalType >= H265NalType::HEVC_TRAIL_N && nalType <= H265NalType::HEVC_CRA_NUT;
221 }
222 
IsIDR(uint8_t nalType)223 bool StartCodeDetectorH265::IsIDR(uint8_t nalType)
224 {
225     return nalType == H265NalType::HEVC_IDR_W_RADL ||
226            nalType == H265NalType::HEVC_IDR_N_LP ||
227            nalType == H265NalType::HEVC_CRA_NUT;
228 }
229 
IsPrefixSEI(uint8_t nalType)230 bool StartCodeDetectorH265::IsPrefixSEI(uint8_t nalType)
231 {
232     return nalType == H265NalType::HEVC_PREFIX_SEI_NUT;
233 }
234 
GetNalType(uint8_t,uint8_t secondByte)235 uint8_t StartCodeDetectorH266::GetNalType(uint8_t, uint8_t secondByte)
236 {
237     return (secondByte & 0b1111'1000) >> 3; // 3: The higher 5 bits of this byte indicate the NalType.
238 }
239 
IsPPS(uint8_t nalType)240 bool StartCodeDetectorH266::IsPPS(uint8_t nalType)
241 {
242     return nalType == H266NalType::VVC_PPS_NUT;
243 }
244 
IsVCL(uint8_t nalType)245 bool StartCodeDetectorH266::IsVCL(uint8_t nalType)
246 {
247     return nalType >= H266NalType::VVC_TRAIL_NUT && nalType <= H266NalType::VVC_RSV_IRAP_11;
248 }
249 
IsIDR(uint8_t nalType)250 bool StartCodeDetectorH266::IsIDR(uint8_t nalType)
251 {
252     return nalType == H266NalType::VVC_IDR_W_RADL ||
253            nalType == H266NalType::VVC_IDR_N_LP ||
254            nalType == H266NalType::VVC_CRA_NUT;
255 ;
256 }
257 
IsPrefixSEI(uint8_t nalType)258 bool StartCodeDetectorH266::IsPrefixSEI(uint8_t nalType)
259 {
260     return nalType == H266NalType::VVC_PREFIX_SEI_NUT;
261 }