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
SetSource(const std::string & path,CodeType type)22 size_t StartCodeDetector::SetSource(const std::string &path, CodeType type)
23 {
24 type_ = type;
25 ifstream ifs(path, ios::binary);
26 if (!ifs.is_open()) {
27 LOGE("cannot open %s", path.c_str());
28 return 0;
29 }
30 size_t fileSize = GetFileSizeInBytes(ifs);
31 unique_ptr<uint8_t[]> buf = make_unique<uint8_t[]>(fileSize);
32 ifs.read(reinterpret_cast<char *>(buf.get()), static_cast<std::streamsize>(fileSize));
33
34 list<pair<size_t, FirstByteInNalu>> posOfFile;
35 const uint8_t *pStart = buf.get();
36 size_t pos = 0;
37 while (pos < fileSize) {
38 auto pFound = search(pStart + pos, pStart + fileSize, begin(START_CODE), end(START_CODE));
39 pos = distance(pStart, pFound);
40 if (pos == fileSize || pos + START_CODE_LEN >= fileSize) { // 没找到或找到的起始码正好在文件末尾
41 break;
42 }
43 posOfFile.emplace_back(pos, pStart[pos + START_CODE_LEN]);
44 pos += START_CODE_LEN;
45 }
46 for (auto it = posOfFile.begin(); it != posOfFile.end(); ++it) {
47 auto nex = next(it);
48 NALUInfo info {
49 .startPos = it->first,
50 .endPos = (nex == posOfFile.end()) ? (fileSize) : (nex->first),
51 .nalType = GetNalType(it->second, type),
52 };
53 nals_.push_back(info);
54 }
55 BuildSampleList();
56 return samples_.size();
57 }
58
BuildSampleList()59 void StartCodeDetector::BuildSampleList()
60 {
61 for (auto it = nals_.begin(); it != nals_.end(); ++it) {
62 Sample sample {
63 .startPos = it->startPos,
64 .endPos = it->endPos,
65 .isCsd = false,
66 .isIdr = false,
67 };
68 if (IsCsd(*it, type_)) {
69 sample.isCsd = true;
70 while (true) {
71 auto nex = next(it);
72 if (nex == nals_.end()) {
73 break;
74 }
75 if (IsCsd(*nex, type_)) {
76 sample.endPos = nex->endPos;
77 it++;
78 } else {
79 break;
80 }
81 }
82 } else if (IsIdr(*it, type_)) {
83 sample.isIdr = true;
84 }
85 if (sample.isCsd) {
86 csdIdxList_.push_back(samples_.size());
87 } else if (sample.isIdr) {
88 idrIdxList_.push_back(samples_.size());
89 }
90 sample.idx = samples_.size();
91 samples_.push_back(sample);
92 }
93 }
94
GetFileSizeInBytes(ifstream & ifs)95 size_t StartCodeDetector::GetFileSizeInBytes(ifstream &ifs)
96 {
97 ifs.seekg(0, ifstream::end);
98 auto len = ifs.tellg();
99 ifs.seekg(0, ifstream::beg);
100 return static_cast<size_t>(len);
101 }
102
GetNalType(uint8_t byte,CodeType type)103 uint8_t StartCodeDetector::GetNalType(uint8_t byte, CodeType type)
104 {
105 switch (type) {
106 case H264: {
107 return byte & 0b0001'1111;
108 }
109 case H265: {
110 return (byte & 0b0111'1110) >> 1;
111 }
112 default: {
113 return 0;
114 }
115 }
116 }
117
IsCsd(const NALUInfo & nalu,CodeType type)118 bool StartCodeDetector::IsCsd(const NALUInfo &nalu, CodeType type)
119 {
120 uint8_t nalType = nalu.nalType;
121 switch (type) {
122 case H264:
123 return nalType == static_cast<uint8_t>(H264NalType::SPS) ||
124 nalType == static_cast<uint8_t>(H264NalType::PPS);
125 case H265:
126 return nalType == static_cast<uint8_t>(H265NalType::HEVC_VPS_NUT) ||
127 nalType == static_cast<uint8_t>(H265NalType::HEVC_SPS_NUT) ||
128 nalType == static_cast<uint8_t>(H265NalType::HEVC_PPS_NUT);
129 default:
130 return false;
131 }
132 }
133
IsIdr(const NALUInfo & nalu,CodeType type)134 bool StartCodeDetector::IsIdr(const NALUInfo &nalu, CodeType type)
135 {
136 uint8_t nalType = nalu.nalType;
137 switch (type) {
138 case H264:
139 return nalType == static_cast<uint8_t>(H264NalType::IDR);
140 case H265:
141 return nalType == static_cast<uint8_t>(H265NalType::HEVC_IDR_W_RADL) ||
142 nalType == static_cast<uint8_t>(H265NalType::HEVC_IDR_N_LP) ||
143 nalType == static_cast<uint8_t>(H265NalType::HEVC_CRA_NUT);
144 default:
145 return false;
146 }
147 }
148
SeekTo(size_t sampleIdx)149 bool StartCodeDetector::SeekTo(size_t sampleIdx)
150 {
151 if (sampleIdx >= samples_.size()) {
152 return false;
153 }
154
155 auto idrIter = find_if(idrIdxList_.rbegin(), idrIdxList_.rend(), [sampleIdx](size_t idrIdx) {
156 return idrIdx <= sampleIdx;
157 });
158 if (idrIter == idrIdxList_.rend()) {
159 return false;
160 }
161 size_t idrIdx = *idrIter;
162
163 auto csdIter = find_if(csdIdxList_.rbegin(), csdIdxList_.rend(), [idrIdx](size_t csdIdx) {
164 return csdIdx < idrIdx;
165 });
166 if (csdIter == csdIdxList_.rend()) {
167 return false;
168 }
169 size_t csdIdx = *csdIter;
170 waitingCsd_ = csdIdx;
171 nextSampleIdx_ = idrIdx;
172 LOGI("csd idx=%zu, idr idx=%zu, target sample idx=%zu", csdIdx, idrIdx, sampleIdx);
173 return true;
174 }
175
PeekNextSample()176 std::optional<Sample> StartCodeDetector::PeekNextSample()
177 {
178 if (waitingCsd_.has_value()) {
179 return samples_[waitingCsd_.value()];
180 }
181 if (nextSampleIdx_ >= samples_.size()) {
182 return std::nullopt;
183 }
184 return samples_[nextSampleIdx_];
185 }
186
MoveToNext()187 void StartCodeDetector::MoveToNext()
188 {
189 if (waitingCsd_.has_value()) {
190 waitingCsd_ = nullopt;
191 return;
192 }
193 nextSampleIdx_++;
194 }