xref: /aosp_15_r20/external/webrtc/modules/video_coding/h264_sps_pps_tracker.cc (revision d9f758449e529ab9291ac668be2861e7a55c2422)
1 /*
2  *  Copyright (c) 2016 The WebRTC project authors. All Rights Reserved.
3  *
4  *  Use of this source code is governed by a BSD-style license
5  *  that can be found in the LICENSE file in the root of the source
6  *  tree. An additional intellectual property rights grant can be found
7  *  in the file PATENTS.  All contributing project authors may
8  *  be found in the AUTHORS file in the root of the source tree.
9  */
10 
11 #include "modules/video_coding/h264_sps_pps_tracker.h"
12 
13 #include <memory>
14 #include <string>
15 #include <utility>
16 
17 #include "absl/types/variant.h"
18 #include "common_video/h264/h264_common.h"
19 #include "common_video/h264/pps_parser.h"
20 #include "common_video/h264/sps_parser.h"
21 #include "modules/video_coding/codecs/h264/include/h264_globals.h"
22 #include "rtc_base/checks.h"
23 #include "rtc_base/logging.h"
24 
25 namespace webrtc {
26 namespace video_coding {
27 
28 namespace {
29 const uint8_t start_code_h264[] = {0, 0, 0, 1};
30 }  // namespace
31 
32 H264SpsPpsTracker::H264SpsPpsTracker() = default;
33 H264SpsPpsTracker::~H264SpsPpsTracker() = default;
34 
35 H264SpsPpsTracker::PpsInfo::PpsInfo() = default;
36 H264SpsPpsTracker::PpsInfo::PpsInfo(PpsInfo&& rhs) = default;
37 H264SpsPpsTracker::PpsInfo& H264SpsPpsTracker::PpsInfo::operator=(
38     PpsInfo&& rhs) = default;
39 H264SpsPpsTracker::PpsInfo::~PpsInfo() = default;
40 
41 H264SpsPpsTracker::SpsInfo::SpsInfo() = default;
42 H264SpsPpsTracker::SpsInfo::SpsInfo(SpsInfo&& rhs) = default;
43 H264SpsPpsTracker::SpsInfo& H264SpsPpsTracker::SpsInfo::operator=(
44     SpsInfo&& rhs) = default;
45 H264SpsPpsTracker::SpsInfo::~SpsInfo() = default;
46 
CopyAndFixBitstream(rtc::ArrayView<const uint8_t> bitstream,RTPVideoHeader * video_header)47 H264SpsPpsTracker::FixedBitstream H264SpsPpsTracker::CopyAndFixBitstream(
48     rtc::ArrayView<const uint8_t> bitstream,
49     RTPVideoHeader* video_header) {
50   RTC_DCHECK(video_header);
51   RTC_DCHECK(video_header->codec == kVideoCodecH264);
52   RTC_DCHECK_GT(bitstream.size(), 0);
53 
54   auto& h264_header =
55       absl::get<RTPVideoHeaderH264>(video_header->video_type_header);
56 
57   bool append_sps_pps = false;
58   auto sps = sps_data_.end();
59   auto pps = pps_data_.end();
60 
61   for (size_t i = 0; i < h264_header.nalus_length; ++i) {
62     const NaluInfo& nalu = h264_header.nalus[i];
63     switch (nalu.type) {
64       case H264::NaluType::kSps: {
65         SpsInfo& sps_info = sps_data_[nalu.sps_id];
66         sps_info.width = video_header->width;
67         sps_info.height = video_header->height;
68         break;
69       }
70       case H264::NaluType::kPps: {
71         pps_data_[nalu.pps_id].sps_id = nalu.sps_id;
72         break;
73       }
74       case H264::NaluType::kIdr: {
75         // If this is the first packet of an IDR, make sure we have the required
76         // SPS/PPS and also calculate how much extra space we need in the buffer
77         // to prepend the SPS/PPS to the bitstream with start codes.
78         if (video_header->is_first_packet_in_frame) {
79           if (nalu.pps_id == -1) {
80             RTC_LOG(LS_WARNING) << "No PPS id in IDR nalu.";
81             return {kRequestKeyframe};
82           }
83 
84           pps = pps_data_.find(nalu.pps_id);
85           if (pps == pps_data_.end()) {
86             RTC_LOG(LS_WARNING)
87                 << "No PPS with id << " << nalu.pps_id << " received";
88             return {kRequestKeyframe};
89           }
90 
91           sps = sps_data_.find(pps->second.sps_id);
92           if (sps == sps_data_.end()) {
93             RTC_LOG(LS_WARNING)
94                 << "No SPS with id << " << pps->second.sps_id << " received";
95             return {kRequestKeyframe};
96           }
97 
98           // Since the first packet of every keyframe should have its width and
99           // height set we set it here in the case of it being supplied out of
100           // band.
101           video_header->width = sps->second.width;
102           video_header->height = sps->second.height;
103 
104           // If the SPS/PPS was supplied out of band then we will have saved
105           // the actual bitstream in `data`.
106           if (sps->second.data && pps->second.data) {
107             RTC_DCHECK_GT(sps->second.size, 0);
108             RTC_DCHECK_GT(pps->second.size, 0);
109             append_sps_pps = true;
110           }
111         }
112         break;
113       }
114       default:
115         break;
116     }
117   }
118 
119   RTC_CHECK(!append_sps_pps ||
120             (sps != sps_data_.end() && pps != pps_data_.end()));
121 
122   // Calculate how much space we need for the rest of the bitstream.
123   size_t required_size = 0;
124 
125   if (append_sps_pps) {
126     required_size += sps->second.size + sizeof(start_code_h264);
127     required_size += pps->second.size + sizeof(start_code_h264);
128   }
129 
130   if (h264_header.packetization_type == kH264StapA) {
131     const uint8_t* nalu_ptr = bitstream.data() + 1;
132     while (nalu_ptr < bitstream.data() + bitstream.size() - 1) {
133       RTC_DCHECK(video_header->is_first_packet_in_frame);
134       required_size += sizeof(start_code_h264);
135 
136       // The first two bytes describe the length of a segment.
137       uint16_t segment_length = nalu_ptr[0] << 8 | nalu_ptr[1];
138       nalu_ptr += 2;
139 
140       required_size += segment_length;
141       nalu_ptr += segment_length;
142     }
143   } else {
144     if (h264_header.nalus_length > 0) {
145       required_size += sizeof(start_code_h264);
146     }
147     required_size += bitstream.size();
148   }
149 
150   // Then we copy to the new buffer.
151   H264SpsPpsTracker::FixedBitstream fixed;
152   fixed.bitstream.EnsureCapacity(required_size);
153 
154   if (append_sps_pps) {
155     // Insert SPS.
156     fixed.bitstream.AppendData(start_code_h264);
157     fixed.bitstream.AppendData(sps->second.data.get(), sps->second.size);
158 
159     // Insert PPS.
160     fixed.bitstream.AppendData(start_code_h264);
161     fixed.bitstream.AppendData(pps->second.data.get(), pps->second.size);
162 
163     // Update codec header to reflect the newly added SPS and PPS.
164     NaluInfo sps_info;
165     sps_info.type = H264::NaluType::kSps;
166     sps_info.sps_id = sps->first;
167     sps_info.pps_id = -1;
168     NaluInfo pps_info;
169     pps_info.type = H264::NaluType::kPps;
170     pps_info.sps_id = sps->first;
171     pps_info.pps_id = pps->first;
172     if (h264_header.nalus_length + 2 <= kMaxNalusPerPacket) {
173       h264_header.nalus[h264_header.nalus_length++] = sps_info;
174       h264_header.nalus[h264_header.nalus_length++] = pps_info;
175     } else {
176       RTC_LOG(LS_WARNING) << "Not enough space in H.264 codec header to insert "
177                              "SPS/PPS provided out-of-band.";
178     }
179   }
180 
181   // Copy the rest of the bitstream and insert start codes.
182   if (h264_header.packetization_type == kH264StapA) {
183     const uint8_t* nalu_ptr = bitstream.data() + 1;
184     while (nalu_ptr < bitstream.data() + bitstream.size() - 1) {
185       fixed.bitstream.AppendData(start_code_h264);
186 
187       // The first two bytes describe the length of a segment.
188       uint16_t segment_length = nalu_ptr[0] << 8 | nalu_ptr[1];
189       nalu_ptr += 2;
190 
191       size_t copy_end = nalu_ptr - bitstream.data() + segment_length;
192       if (copy_end > bitstream.size()) {
193         return {kDrop};
194       }
195 
196       fixed.bitstream.AppendData(nalu_ptr, segment_length);
197       nalu_ptr += segment_length;
198     }
199   } else {
200     if (h264_header.nalus_length > 0) {
201       fixed.bitstream.AppendData(start_code_h264);
202     }
203     fixed.bitstream.AppendData(bitstream.data(), bitstream.size());
204   }
205 
206   fixed.action = kInsert;
207   return fixed;
208 }
209 
InsertSpsPpsNalus(const std::vector<uint8_t> & sps,const std::vector<uint8_t> & pps)210 void H264SpsPpsTracker::InsertSpsPpsNalus(const std::vector<uint8_t>& sps,
211                                           const std::vector<uint8_t>& pps) {
212   constexpr size_t kNaluHeaderOffset = 1;
213   if (sps.size() < kNaluHeaderOffset) {
214     RTC_LOG(LS_WARNING) << "SPS size  " << sps.size() << " is smaller than "
215                         << kNaluHeaderOffset;
216     return;
217   }
218   if ((sps[0] & 0x1f) != H264::NaluType::kSps) {
219     RTC_LOG(LS_WARNING) << "SPS Nalu header missing";
220     return;
221   }
222   if (pps.size() < kNaluHeaderOffset) {
223     RTC_LOG(LS_WARNING) << "PPS size  " << pps.size() << " is smaller than "
224                         << kNaluHeaderOffset;
225     return;
226   }
227   if ((pps[0] & 0x1f) != H264::NaluType::kPps) {
228     RTC_LOG(LS_WARNING) << "SPS Nalu header missing";
229     return;
230   }
231   absl::optional<SpsParser::SpsState> parsed_sps = SpsParser::ParseSps(
232       sps.data() + kNaluHeaderOffset, sps.size() - kNaluHeaderOffset);
233   absl::optional<PpsParser::PpsState> parsed_pps = PpsParser::ParsePps(
234       pps.data() + kNaluHeaderOffset, pps.size() - kNaluHeaderOffset);
235 
236   if (!parsed_sps) {
237     RTC_LOG(LS_WARNING) << "Failed to parse SPS.";
238   }
239 
240   if (!parsed_pps) {
241     RTC_LOG(LS_WARNING) << "Failed to parse PPS.";
242   }
243 
244   if (!parsed_pps || !parsed_sps) {
245     return;
246   }
247 
248   SpsInfo sps_info;
249   sps_info.size = sps.size();
250   sps_info.width = parsed_sps->width;
251   sps_info.height = parsed_sps->height;
252   uint8_t* sps_data = new uint8_t[sps_info.size];
253   memcpy(sps_data, sps.data(), sps_info.size);
254   sps_info.data.reset(sps_data);
255   sps_data_[parsed_sps->id] = std::move(sps_info);
256 
257   PpsInfo pps_info;
258   pps_info.size = pps.size();
259   pps_info.sps_id = parsed_pps->sps_id;
260   uint8_t* pps_data = new uint8_t[pps_info.size];
261   memcpy(pps_data, pps.data(), pps_info.size);
262   pps_info.data.reset(pps_data);
263   pps_data_[parsed_pps->id] = std::move(pps_info);
264 
265   RTC_LOG(LS_INFO) << "Inserted SPS id " << parsed_sps->id << " and PPS id "
266                    << parsed_pps->id << " (referencing SPS "
267                    << parsed_pps->sps_id << ")";
268 }
269 
270 }  // namespace video_coding
271 }  // namespace webrtc
272