1 /*
2 mkvextract -- extract tracks from Matroska files into other files
3
4 Distributed under the GPL v2
5 see the file COPYING for details
6 or visit https://www.gnu.org/licenses/old-licenses/gpl-2.0.html
7
8 extracts tracks from Matroska files into other files
9
10 */
11
12 #include "common/common_pch.h"
13
14 #include "common/endian.h"
15 #include "common/hacks.h"
16 #include "common/hevc/util.h"
17 #include "common/list_utils.h"
18 #include "extract/xtr_hevc.h"
19
xtr_hevc_c(std::string const & codec_id,int64_t tid,track_spec_t & tspec)20 xtr_hevc_c::xtr_hevc_c(std::string const &codec_id,
21 int64_t tid,
22 track_spec_t &tspec)
23 : xtr_avc_c{codec_id, tid, tspec}
24 , m_normalize_parameter_sets{!mtx::hacks::is_engaged(mtx::hacks::DONT_NORMALIZE_PARAMETER_SETS)}
25 {
26 }
27
28 void
create_file(xtr_base_c * master,libmatroska::KaxTrackEntry & track)29 xtr_hevc_c::create_file(xtr_base_c *master,
30 libmatroska::KaxTrackEntry &track) {
31 xtr_base_c::create_file(master, track);
32
33 auto priv = FindChild<libmatroska::KaxCodecPrivate>(&track);
34 if (!priv)
35 mxerror(fmt::format(Y("Track {0} with the CodecID '{1}' is missing the \"codec private\" element and cannot be extracted.\n"), m_tid, m_codec_id));
36
37 m_decoded_codec_private = decode_codec_private(priv);
38 if (m_decoded_codec_private->get_size() < 23)
39 mxerror(fmt::format(Y("Track {0} CodecPrivate is too small.\n"), m_tid));
40
41 m_parser.set_normalize_parameter_sets(m_normalize_parameter_sets);
42 m_parser.set_configuration_record(m_decoded_codec_private);
43
44 m_nal_size_size = 1 + (m_decoded_codec_private->get_buffer()[21] & 3);
45 }
46
47 void
unwrap_write_hevcc(bool skip_prefix_sei)48 xtr_hevc_c::unwrap_write_hevcc(bool skip_prefix_sei) {
49 // Parameter sets in this order: vps, sps, pps, sei
50 auto buf = m_decoded_codec_private->get_buffer();
51 auto num_parameter_sets = static_cast<unsigned int>(buf[22]);
52 auto pos = static_cast<size_t>(23);
53 auto const priv_size = m_decoded_codec_private->get_size();
54
55 mxdebug_if(m_debug, fmt::format("unwrap_write_hevcc: nal_size_size {0} num_parameter_sets {1}\n", m_nal_size_size, num_parameter_sets));
56
57 while (num_parameter_sets && (priv_size > (pos + 3))) {
58 auto nal_unit_count = get_uint16_be(&buf[pos + 1]);
59 pos += 3;
60
61 mxdebug_if(m_debug, fmt::format("unwrap_write_hevcc: num_parameter_sets {0} nal_unit_count {1} pos {2}\n", num_parameter_sets, nal_unit_count, pos - 3));
62
63 while (nal_unit_count && (priv_size > pos)) {
64 if ((pos + 2 + 1) > priv_size)
65 return;
66
67 auto nal_size = get_uint_be(&buf[pos], 2);
68 auto nal_unit_type = (buf[pos + 2] >> 1) & 0x3f;
69 auto skip_this_nal = skip_prefix_sei && (mtx::hevc::NALU_TYPE_PREFIX_SEI == nal_unit_type);
70
71 mxdebug_if(m_debug, fmt::format("unwrap_write_hevcc: pos {0} nal_size {1} type 0x{2:02x} skip_this_nal {3}\n", pos, nal_size, static_cast<unsigned int>(nal_unit_type), skip_this_nal));
72
73 if (skip_this_nal)
74 pos += nal_size + 2;
75
76 else if (!write_nal(buf, pos, priv_size, 2))
77 return;
78
79 --nal_unit_count;
80 --num_parameter_sets;
81 }
82 }
83 }
84
85 bool
write_nal(binary * data,size_t & pos,size_t data_size,size_t write_nal_size_size)86 xtr_hevc_c::write_nal(binary *data,
87 size_t &pos,
88 size_t data_size,
89 size_t write_nal_size_size) {
90 if ((pos + write_nal_size_size) > data_size)
91 return false;
92
93 auto nal_size = get_uint_be(&data[pos], write_nal_size_size);
94 pos += write_nal_size_size;
95
96 if ((pos + nal_size) > data_size) {
97 mxwarn(fmt::format(Y("Track {0}: NAL too big. Size according to header field: {1}, available bytes in packet: {2}. This NAL is defect and will be skipped.\n"), m_tid, nal_size, data_size - pos));
98 return false;
99 }
100
101 auto nal_unit_type = (data[pos] >> 1) & 0x3f;
102 auto start_code_size = m_first_nalu || mtx::included_in(nal_unit_type, mtx::hevc::NALU_TYPE_VIDEO_PARAM, mtx::hevc::NALU_TYPE_SEQ_PARAM, mtx::hevc::NALU_TYPE_PIC_PARAM) ? 4 : 3;
103 m_first_nalu = false;
104
105 m_out->write(ms_start_code + (4 - start_code_size), start_code_size);
106 m_out->write(data + pos, nal_size);
107
108 pos += nal_size;
109
110 return true;
111 }
112
113 void
handle_frame(xtr_frame_t & f)114 xtr_hevc_c::handle_frame(xtr_frame_t &f) {
115 if ( m_first_nalu
116 && !m_normalize_parameter_sets
117 && (f.frame->get_size() > static_cast<size_t>(m_nal_size_size))) {
118 auto nal_unit_type = (f.frame->get_buffer()[m_nal_size_size + 1] >> 1) & 0x3f;
119 if (!mtx::included_in(nal_unit_type, mtx::hevc::NALU_TYPE_VIDEO_PARAM, mtx::hevc::NALU_TYPE_SEQ_PARAM, mtx::hevc::NALU_TYPE_PIC_PARAM))
120 unwrap_write_hevcc(nal_unit_type == mtx::hevc::NALU_TYPE_PREFIX_SEI);
121 }
122
123 if (f.keyframe)
124 m_parser.set_next_i_slice_is_key_frame();
125
126 m_parser.add_bytes_framed(f.frame, m_nal_size_size);
127
128 flush_frames();
129 }
130
131 void
flush_frames()132 xtr_hevc_c::flush_frames() {
133 while (m_parser.frame_available()) {
134 std::size_t pos{};
135
136 auto frame = m_parser.get_frame();
137 auto buffer = frame.m_data->get_buffer();
138 auto size = frame.m_data->get_size();
139
140 while (pos < size)
141 write_nal(buffer, pos, size, m_nal_size_size);
142 }
143 }
144
145 void
finish_track()146 xtr_hevc_c::finish_track() {
147 m_parser.flush();
148 flush_frames();
149 }
150