diff --git a/.dockerignore b/.dockerignore index 1f261adf30..cfc34e5848 100644 --- a/.dockerignore +++ b/.dockerignore @@ -28,10 +28,8 @@ chffr/app2 chffr/backend/env selfdrive/nav selfdrive/baseui -chffr/lib/vidindex/vidindex selfdrive/test/simulator2 **/cache_data -xx/chffr/lib/vidindex/vidindex xx/plus xx/community xx/projects diff --git a/tools/lib/framereader.py b/tools/lib/framereader.py index 06fee6857d..e0b24963ca 100644 --- a/tools/lib/framereader.py +++ b/tools/lib/framereader.py @@ -14,6 +14,7 @@ from lru import LRU import _io from openpilot.tools.lib.cache import cache_path_for_file_path, DEFAULT_CACHE_DIR from openpilot.tools.lib.exceptions import DataUnreadableError +from openpilot.tools.lib.vidindex import hevc_index from openpilot.common.file_helpers import atomic_write_in_dir from openpilot.tools.lib.filereader import FileReader @@ -75,31 +76,6 @@ def ffprobe(fn, fmt=None): return json.loads(ffprobe_output) -def vidindex(fn, typ): - vidindex_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "vidindex") - vidindex = os.path.join(vidindex_dir, "vidindex") - - subprocess.check_call(["make"], cwd=vidindex_dir, stdout=subprocess.DEVNULL) - - with tempfile.NamedTemporaryFile() as prefix_f, \ - tempfile.NamedTemporaryFile() as index_f: - try: - subprocess.check_call([vidindex, typ, fn, prefix_f.name, index_f.name]) - except subprocess.CalledProcessError as e: - raise DataUnreadableError(f"vidindex failed on file {fn}") from e - with open(index_f.name, "rb") as f: - index = f.read() - with open(prefix_f.name, "rb") as f: - prefix = f.read() - - index = np.frombuffer(index, np.uint32).reshape(-1, 2) - - assert index[-1, 0] == 0xFFFFFFFF - assert index[-1, 1] == os.path.getsize(fn) - - return index, prefix - - def cache_fn(func): @wraps(func) def cache_inner(fn, *args, **kwargs): @@ -114,7 +90,6 @@ def cache_fn(func): cache_value = pickle.load(cache_file) else: cache_value = func(fn, *args, **kwargs) - if cache_path: with atomic_write_in_dir(cache_path, mode="wb", overwrite=True) as cache_file: pickle.dump(cache_value, cache_file, -1) @@ -125,13 +100,13 @@ def cache_fn(func): @cache_fn -def index_stream(fn, typ): - assert typ in ("hevc", ) +def index_stream(fn, ft): + if ft != FrameType.h265_stream: + raise NotImplementedError("Only h265 supported") - with FileReader(fn) as f: - assert os.path.exists(f.name), fn - index, prefix = vidindex(f.name, typ) - probe = ffprobe(f.name, typ) + frame_types, dat_len, prefix = hevc_index(fn) + index = np.array(frame_types + [(0xFFFFFFFF, dat_len)], dtype=np.uint32) + probe = ffprobe(fn, "hevc") return { 'index': index, @@ -140,42 +115,8 @@ def index_stream(fn, typ): } -def index_videos(camera_paths, cache_dir=DEFAULT_CACHE_DIR): - """Requires that paths in camera_paths are contiguous and of the same type.""" - if len(camera_paths) < 1: - raise ValueError("must provide at least one video to index") - - frame_type = fingerprint_video(camera_paths[0]) - for fn in camera_paths: - index_video(fn, frame_type, cache_dir) - - -def index_video(fn, frame_type=None, cache_dir=DEFAULT_CACHE_DIR): - cache_path = cache_path_for_file_path(fn, cache_dir) - - if os.path.exists(cache_path): - return - - if frame_type is None: - frame_type = fingerprint_video(fn[0]) - - if frame_type == FrameType.h265_stream: - index_stream(fn, "hevc", cache_dir=cache_dir) - else: - raise NotImplementedError("Only h265 supported") - - def get_video_index(fn, frame_type, cache_dir=DEFAULT_CACHE_DIR): - cache_path = cache_path_for_file_path(fn, cache_dir) - - if not os.path.exists(cache_path): - index_video(fn, frame_type, cache_dir) - - if not os.path.exists(cache_path): - return None - with open(cache_path, "rb") as cache_file: - return pickle.load(cache_file) - + return index_stream(fn, frame_type, cache_dir=cache_dir) def read_file_check_size(f, sz, cookie): buff = bytearray(sz) diff --git a/tools/lib/vidindex.py b/tools/lib/vidindex.py new file mode 100755 index 0000000000..9404d035ba --- /dev/null +++ b/tools/lib/vidindex.py @@ -0,0 +1,301 @@ +#!/usr/bin/env python3 +import argparse +import struct +from enum import IntEnum +from typing import Tuple + +from openpilot.tools.lib.filereader import FileReader + +# H.265 specification +# https://www.itu.int/rec/dologin_pub.asp?lang=e&id=T-REC-H.265-201802-S!!PDF-E&type=items + +NAL_UNIT_START_CODE = b"\x00\x00\x01" +NAL_UNIT_START_CODE_SIZE = len(NAL_UNIT_START_CODE) +NAL_UNIT_HEADER_SIZE = 2 + +class HevcNalUnitType(IntEnum): + TRAIL_N = 0 # RBSP structure: slice_segment_layer_rbsp( ) + TRAIL_R = 1 # RBSP structure: slice_segment_layer_rbsp( ) + TSA_N = 2 # RBSP structure: slice_segment_layer_rbsp( ) + TSA_R = 3 # RBSP structure: slice_segment_layer_rbsp( ) + STSA_N = 4 # RBSP structure: slice_segment_layer_rbsp( ) + STSA_R = 5 # RBSP structure: slice_segment_layer_rbsp( ) + RADL_N = 6 # RBSP structure: slice_segment_layer_rbsp( ) + RADL_R = 7 # RBSP structure: slice_segment_layer_rbsp( ) + RASL_N = 8 # RBSP structure: slice_segment_layer_rbsp( ) + RASL_R = 9 # RBSP structure: slice_segment_layer_rbsp( ) + RSV_VCL_N10 = 10 + RSV_VCL_R11 = 11 + RSV_VCL_N12 = 12 + RSV_VCL_R13 = 13 + RSV_VCL_N14 = 14 + RSV_VCL_R15 = 15 + BLA_W_LP = 16 # RBSP structure: slice_segment_layer_rbsp( ) + BLA_W_RADL = 17 # RBSP structure: slice_segment_layer_rbsp( ) + BLA_N_LP = 18 # RBSP structure: slice_segment_layer_rbsp( ) + IDR_W_RADL = 19 # RBSP structure: slice_segment_layer_rbsp( ) + IDR_N_LP = 20 # RBSP structure: slice_segment_layer_rbsp( ) + CRA_NUT = 21 # RBSP structure: slice_segment_layer_rbsp( ) + RSV_IRAP_VCL22 = 22 + RSV_IRAP_VCL23 = 23 + RSV_VCL24 = 24 + RSV_VCL25 = 25 + RSV_VCL26 = 26 + RSV_VCL27 = 27 + RSV_VCL28 = 28 + RSV_VCL29 = 29 + RSV_VCL30 = 30 + RSV_VCL31 = 31 + VPS_NUT = 32 # RBSP structure: video_parameter_set_rbsp( ) + SPS_NUT = 33 # RBSP structure: seq_parameter_set_rbsp( ) + PPS_NUT = 34 # RBSP structure: pic_parameter_set_rbsp( ) + AUD_NUT = 35 + EOS_NUT = 36 + EOB_NUT = 37 + FD_NUT = 38 + PREFIX_SEI_NUT = 39 + SUFFIX_SEI_NUT = 40 + RSV_NVCL41 = 41 + RSV_NVCL42 = 42 + RSV_NVCL43 = 43 + RSV_NVCL44 = 44 + RSV_NVCL45 = 45 + RSV_NVCL46 = 46 + RSV_NVCL47 = 47 + UNSPEC48 = 48 + UNSPEC49 = 49 + UNSPEC50 = 50 + UNSPEC51 = 51 + UNSPEC52 = 52 + UNSPEC53 = 53 + UNSPEC54 = 54 + UNSPEC55 = 55 + UNSPEC56 = 56 + UNSPEC57 = 57 + UNSPEC58 = 58 + UNSPEC59 = 59 + UNSPEC60 = 60 + UNSPEC61 = 61 + UNSPEC62 = 62 + UNSPEC63 = 63 + +# B.2.2 Byte stream NAL unit semantics +# - The nal_unit_type within the nal_unit( ) syntax structure is equal to VPS_NUT, SPS_NUT or PPS_NUT. +# - The byte stream NAL unit syntax structure contains the first NAL unit of an access unit in decoding +# order, as specified in clause 7.4.2.4.4. +HEVC_PARAMETER_SET_NAL_UNITS = ( + HevcNalUnitType.VPS_NUT, + HevcNalUnitType.SPS_NUT, + HevcNalUnitType.PPS_NUT, +) + +# 3.29 coded slice segment NAL unit: A NAL unit that has nal_unit_type in the range of TRAIL_N to RASL_R, +# inclusive, or in the range of BLA_W_LP to RSV_IRAP_VCL23, inclusive, which indicates that the NAL unit +# contains a coded slice segment +HEVC_CODED_SLICE_SEGMENT_NAL_UNITS = ( + HevcNalUnitType.TRAIL_N, + HevcNalUnitType.TRAIL_R, + HevcNalUnitType.TSA_N, + HevcNalUnitType.TSA_R, + HevcNalUnitType.STSA_N, + HevcNalUnitType.STSA_R, + HevcNalUnitType.RADL_N, + HevcNalUnitType.RADL_R, + HevcNalUnitType.RASL_N, + HevcNalUnitType.RASL_R, + HevcNalUnitType.BLA_W_LP, + HevcNalUnitType.BLA_W_RADL, + HevcNalUnitType.BLA_N_LP, + HevcNalUnitType.IDR_W_RADL, + HevcNalUnitType.IDR_N_LP, + HevcNalUnitType.CRA_NUT, +) + +class VideoFileInvalid(Exception): + pass + +def get_ue(dat: bytes, start_idx: int, skip_bits: int) -> Tuple[int, int]: + prefix_val = 0 + prefix_len = 0 + suffix_val = 0 + suffix_len = 0 + + i = start_idx + while i < len(dat): + j = 7 + while j >= 0: + if skip_bits > 0: + skip_bits -= 1 + elif prefix_val == 0: + prefix_val = (dat[i] >> j) & 1 + prefix_len += 1 + else: + suffix_val = (suffix_val << 1) | ((dat[i] >> j) & 1) + suffix_len += 1 + j -= 1 + + if prefix_val == 1 and prefix_len - 1 == suffix_len: + val = 2**(prefix_len-1) - 1 + suffix_val + size = prefix_len + suffix_len + return val, size + i += 1 + + raise VideoFileInvalid("invalid exponential-golomb code") + +def require_nal_unit_start(dat: bytes, nal_unit_start: int) -> None: + if nal_unit_start >= len(dat): + raise ValueError("start index must be less than data length") + + if nal_unit_start < 1: + raise ValueError("start index must be greater than zero") + + if dat[nal_unit_start-1] != 0x00: + raise VideoFileInvalid("start code must be preceded by 0x00") + + if dat[nal_unit_start:nal_unit_start + NAL_UNIT_START_CODE_SIZE] != NAL_UNIT_START_CODE: + raise VideoFileInvalid("data must begin with start code") + +def get_hevc_nal_unit_length(dat: bytes, nal_unit_start: int) -> int: + try: + pos = dat.index(NAL_UNIT_START_CODE, nal_unit_start + NAL_UNIT_START_CODE_SIZE) + except ValueError: + pos = -1 + + # length of NAL unit is byte count up to next NAL unit start index + nal_unit_len = (pos if pos != -1 else len(dat)) - nal_unit_start + return nal_unit_len + +def get_hevc_nal_unit_type(dat: bytes, nal_unit_start: int) -> HevcNalUnitType: + # 7.3.1.2 NAL unit header syntax + # nal_unit_header( ) { // descriptor + # forbidden_zero_bit f(1) + # nal_unit_type u(6) + # nuh_layer_id u(6) + # nuh_temporal_id_plus1 u(3) + # } + header_start = nal_unit_start + NAL_UNIT_START_CODE_SIZE + nal_unit_header = dat[header_start:header_start + NAL_UNIT_HEADER_SIZE] + if len(nal_unit_header) != 2: + raise VideoFileInvalid("data to short to contain nal unit header") + return HevcNalUnitType((nal_unit_header[0] >> 1) & 0x3F) + +def get_hevc_slice_type(dat: bytes, nal_unit_start: int, nal_unit_type: HevcNalUnitType) -> Tuple[int, bool]: + # 7.3.2.9 Slice segment layer RBSP syntax + # slice_segment_layer_rbsp( ) { + # slice_segment_header( ) + # slice_segment_data( ) + # rbsp_slice_segment_trailing_bits( ) + # } + # ... + # 7.3.6.1 General slice segment header syntax + # slice_segment_header( ) { // descriptor + # first_slice_segment_in_pic_flag u(1) + # if( nal_unit_type >= BLA_W_LP && nal_unit_type <= RSV_IRAP_VCL23 ) + # no_output_of_prior_pics_flag u(1) + # slice_pic_parameter_set_id ue(v) + # if( !first_slice_segment_in_pic_flag ) { + # if( dependent_slice_segments_enabled_flag ) + # dependent_slice_segment_flag u(1) + # slice_segment_address u(v) + # } + # if( !dependent_slice_segment_flag ) { + # for( i = 0; i < num_extra_slice_header_bits; i++ ) + # slice_reserved_flag[ i ] u(1) + # slice_type ue(v) + # ... + + rbsp_start = nal_unit_start + NAL_UNIT_START_CODE_SIZE + NAL_UNIT_HEADER_SIZE + skip_bits = 0 + + # 7.4.7.1 General slice segment header semantics + # first_slice_segment_in_pic_flag equal to 1 specifies that the slice segment is the first slice segment of the picture in + # decoding order. first_slice_segment_in_pic_flag equal to 0 specifies that the slice segment is not the first slice segment + # of the picture in decoding order. + is_first_slice = dat[rbsp_start] >> 7 & 1 == 1 + if not is_first_slice: + # TODO: parse dependent_slice_segment_flag and slice_segment_address and get real slice_type + # for now since we don't use it return -1 for slice_type + return (-1, is_first_slice) + skip_bits += 1 # skip past first_slice_segment_in_pic_flag + + if nal_unit_type >= HevcNalUnitType.BLA_W_LP and nal_unit_type <= HevcNalUnitType.RSV_IRAP_VCL23: + # 7.4.7.1 General slice segment header semantics + # no_output_of_prior_pics_flag affects the output of previously-decoded pictures in the decoded picture buffer after the + # decoding of an IDR or a BLA picture that is not the first picture in the bitstream as specified in Annex C. + skip_bits += 1 # skip past no_output_of_prior_pics_flag + + # 7.4.7.1 General slice segment header semantics + # slice_pic_parameter_set_id specifies the value of pps_pic_parameter_set_id for the PPS in use. + # The value of slice_pic_parameter_set_id shall be in the range of 0 to 63, inclusive. + _, size = get_ue(dat, rbsp_start, skip_bits) + skip_bits += size # skip past slice_pic_parameter_set_id + + # 7.4.3.3.1 General picture parameter set RBSP semanal_unit_lenntics + # num_extra_slice_header_bits specifies the number of extra slice header bits that are present in the slice header RBSP + # for coded pictures referring to the PPS. The value of num_extra_slice_header_bits shall be in the range of 0 to 2, inclusive, + # in bitstreams conforming to this version of this Specification. Other values for num_extra_slice_header_bits are reserved + # for future use by ITU-T | ISO/IEC. However, decoders shall allow num_extra_slice_header_bits to have any value. + # TODO: get from PPS_NUT pic_parameter_set_rbsp( ) for corresponding slice_pic_parameter_set_id + num_extra_slice_header_bits = 0 + skip_bits += num_extra_slice_header_bits + + # 7.4.7.1 General slice segment header semantics + # slice_type specifies the coding type of the slice according to Table 7-7. + # Table 7-7 - Name association to slice_type + # slice_type | Name of slice_type + # 0 | B (B slice) + # 1 | P (P slice) + # 2 | I (I slice) + # unsigned integer 0-th order Exp-Golomb-coded syntax element with the left bit first + slice_type, _ = get_ue(dat, rbsp_start, skip_bits) + if slice_type > 2: + raise VideoFileInvalid("slice_type must be 0, 1, or 2") + return slice_type, is_first_slice + +def hevc_index(hevc_file_name: str, allow_corrupt: bool=False) -> Tuple[list, int, bytes]: + with FileReader(hevc_file_name) as f: + dat = f.read() + + if len(dat) < NAL_UNIT_START_CODE_SIZE + 1: + raise VideoFileInvalid("data is too short") + + prefix_dat = b"" + frame_types = list() + + try: + i = 1 # skip past first byte 0x00 (verified by get_hevc_nal_info) + while i < len(dat): + require_nal_unit_start(dat, i) + nal_unit_len = get_hevc_nal_unit_length(dat, i) + nal_unit_type = get_hevc_nal_unit_type(dat, i) + if nal_unit_type in HEVC_PARAMETER_SET_NAL_UNITS: + prefix_dat += dat[i:i+nal_unit_len] + elif nal_unit_type in HEVC_CODED_SLICE_SEGMENT_NAL_UNITS: + slice_type, is_first_slice = get_hevc_slice_type(dat, i, nal_unit_type) + if is_first_slice: + frame_types.append((slice_type, i)) + i += nal_unit_len + except Exception: + if not allow_corrupt: + raise + + return frame_types, len(dat), prefix_dat + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("input_file", type=str) + parser.add_argument("output_prefix_file", type=str) + parser.add_argument("output_index_file", type=str) + args = parser.parse_args() + + frame_types, dat_len, prefix_dat = hevc_index(args.input_file) + with open(args.output_prefix_file, "wb") as f: + f.write(prefix_dat) + + with open(args.output_index_file, "wb") as f: + for ft, fp in frame_types: + f.write(struct.pack(" -#include - -static const uint32_t BS_MASKS[33] = { - 0, 0x1L, 0x3L, 0x7L, 0xFL, 0x1FL, - 0x3FL, 0x7FL, 0xFFL, 0x1FFL, 0x3FFL, 0x7FFL, - 0xFFFL, 0x1FFFL, 0x3FFFL, 0x7FFFL, 0xFFFFL, 0x1FFFFL, - 0x3FFFFL, 0x7FFFFL, 0xFFFFFL, 0x1FFFFFL, 0x3FFFFFL, 0x7FFFFFL, - 0xFFFFFFL, 0x1FFFFFFL, 0x3FFFFFFL, 0x7FFFFFFL, 0xFFFFFFFL, 0x1FFFFFFFL, - 0x3FFFFFFFL, 0x7FFFFFFFL, 0xFFFFFFFFL}; - -void bs_init(struct bitstream* bs, const uint8_t* buffer, size_t input_size) { - bs->buffer_ptr = buffer; - bs->buffer_end = buffer + input_size; - bs->value = 0; - bs->pos = 0; - bs->shift = 8; - bs->size = input_size * 8; -} - -uint32_t bs_get(struct bitstream* bs, int n) { - if (n > 32) - return 0; - - bs->pos += n; - bs->shift += n; - while (bs->shift > 8) { - if (bs->buffer_ptr < bs->buffer_end) { - bs->value <<= 8; - bs->value |= *bs->buffer_ptr++; - bs->shift -= 8; - } else { - bs_seek(bs, bs->pos - n); - return 0; - // bs->value <<= 8; - // bs->shift -= 8; - } - } - return (bs->value >> (8 - bs->shift)) & BS_MASKS[n]; -} - -void bs_seek(struct bitstream* bs, size_t new_pos) { - bs->pos = (new_pos / 32) * 32; - bs->shift = 8; - bs->value = 0; - bs_get(bs, new_pos % 32); -} - -uint32_t bs_peek(struct bitstream* bs, int n) { - struct bitstream bak = *bs; - return bs_get(&bak, n); -} - -size_t bs_remain(struct bitstream* bs) { - return bs->size - bs->pos; -} - -int bs_eof(struct bitstream* bs) { - return bs_remain(bs) == 0; -} - -uint32_t bs_ue(struct bitstream* bs) { - static const uint8_t exp_golomb_bits[256] = { - 8, 7, 6, 6, 5, 5, 5, 5, 4, 4, 4, 4, 4, 4, 4, 4, 3, 3, 3, 3, 3, 3, 3, 3, - 3, 3, 3, 3, 3, 3, 3, 3, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, - 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, - 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, - }; - uint32_t bits, read = 0; - int bits_left; - uint8_t coded; - int done = 0; - bits = 0; - // we want to read 8 bits at a time - if we don't have 8 bits, - // read what's left, and shift. The exp_golomb_bits calc remains the - // same. - while (!done) { - bits_left = bs_remain(bs); - if (bits_left < 8) { - read = bs_peek(bs, bits_left) << (8 - bits_left); - done = 1; - } else { - read = bs_peek(bs, 8); - if (read == 0) { - bs_get(bs, 8); - bits += 8; - } else { - done = 1; - } - } - } - coded = exp_golomb_bits[read]; - bs_get(bs, coded); - bits += coded; - - // printf("ue - bits %d\n", bits); - return bs_get(bs, bits + 1) - 1; -} - -int32_t bs_se(struct bitstream* bs) { - uint32_t ret; - ret = bs_ue(bs); - if ((ret & 0x1) == 0) { - ret >>= 1; - int32_t temp = 0 - ret; - return temp; - } - return (ret + 1) >> 1; -} diff --git a/tools/lib/vidindex/bitstream.h b/tools/lib/vidindex/bitstream.h deleted file mode 100644 index 0f538a59ab..0000000000 --- a/tools/lib/vidindex/bitstream.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef bitstream_H -#define bitstream_H - - -#include -#include - -struct bitstream { - const uint8_t *buffer_ptr; - const uint8_t *buffer_end; - uint64_t value; - uint32_t pos; - uint32_t shift; - size_t size; -}; - -void bs_init(struct bitstream *bs, const uint8_t *buffer, size_t input_size); -void bs_seek(struct bitstream *bs, size_t new_pos); -uint32_t bs_get(struct bitstream *bs, int n); -uint32_t bs_peek(struct bitstream *bs, int n); -size_t bs_remain(struct bitstream *bs); -int bs_eof(struct bitstream *bs); -uint32_t bs_ue(struct bitstream *bs); -int32_t bs_se(struct bitstream *bs); - -#endif diff --git a/tools/lib/vidindex/vidindex.c b/tools/lib/vidindex/vidindex.c deleted file mode 100644 index 1d7da2e193..0000000000 --- a/tools/lib/vidindex/vidindex.c +++ /dev/null @@ -1,307 +0,0 @@ -#include -#include -#include -#include -#include - -#include -#include -#include -#include - -#include "./bitstream.h" - -#define START_CODE 0x000001 - -static uint32_t read24be(const uint8_t* ptr) { - return (ptr[0] << 16) | (ptr[1] << 8) | ptr[2]; -} -static void write32le(FILE *of, uint32_t v) { - uint8_t va[4] = { - v & 0xff, (v >> 8) & 0xff, (v >> 16) & 0xff, (v >> 24) & 0xff - }; - fwrite(va, 1, sizeof(va), of); -} - -// Table 7-1 -enum hevc_nal_type { - HEVC_NAL_TYPE_TRAIL_N = 0, - HEVC_NAL_TYPE_TRAIL_R = 1, - HEVC_NAL_TYPE_TSA_N = 2, - HEVC_NAL_TYPE_TSA_R = 3, - HEVC_NAL_TYPE_STSA_N = 4, - HEVC_NAL_TYPE_STSA_R = 5, - HEVC_NAL_TYPE_RADL_N = 6, - HEVC_NAL_TYPE_RADL_R = 7, - HEVC_NAL_TYPE_RASL_N = 8, - HEVC_NAL_TYPE_RASL_R = 9, - HEVC_NAL_TYPE_BLA_W_LP = 16, - HEVC_NAL_TYPE_BLA_W_RADL = 17, - HEVC_NAL_TYPE_BLA_N_LP = 18, - HEVC_NAL_TYPE_IDR_W_RADL = 19, - HEVC_NAL_TYPE_IDR_N_LP = 20, - HEVC_NAL_TYPE_CRA_NUT = 21, - HEVC_NAL_TYPE_RSV_IRAP_VCL23 = 23, - HEVC_NAL_TYPE_VPS_NUT = 32, - HEVC_NAL_TYPE_SPS_NUT = 33, - HEVC_NAL_TYPE_PPS_NUT = 34, - HEVC_NAL_TYPE_AUD_NUT = 35, - HEVC_NAL_TYPE_EOS_NUT = 36, - HEVC_NAL_TYPE_EOB_NUT = 37, - HEVC_NAL_TYPE_FD_NUT = 38, - HEVC_NAL_TYPE_PREFIX_SEI_NUT = 39, - HEVC_NAL_TYPE_SUFFIX_SEI_NUT = 40, -}; - -// Table 7-7 -enum hevc_slice_type { - HEVC_SLICE_B = 0, - HEVC_SLICE_P = 1, - HEVC_SLICE_I = 2, -}; - -static void hevc_index(const uint8_t *data, size_t file_size, FILE *of_prefix, FILE *of_index) { - const uint8_t* ptr = data; - const uint8_t* ptr_end = data + file_size; - - assert(ptr[0] == 0); - ptr++; - assert(read24be(ptr) == START_CODE); - - // pps. ignore for now - uint32_t num_extra_slice_header_bits = 0; - uint32_t dependent_slice_segments_enabled_flag = 0; - - while (ptr < ptr_end) { - const uint8_t* next = ptr+1; - for (; next < ptr_end-4; next++) { - if (read24be(next) == START_CODE) break; - } - size_t nal_size = next - ptr; - if (nal_size < 6) { - break; - } - - { - struct bitstream bs = {0}; - bs_init(&bs, ptr, nal_size); - - uint32_t start_code = bs_get(&bs, 24); - assert(start_code == 0x000001); - - // nal_unit_header - uint32_t forbidden_zero_bit = bs_get(&bs, 1); - uint32_t nal_unit_type = bs_get(&bs, 6); - uint32_t nuh_layer_id = bs_get(&bs, 6); - uint32_t nuh_temporal_id_plus1 = bs_get(&bs, 3); - - // if (nal_unit_type != 1) printf("%3d -- %3d %10d %lu\n", nal_unit_type, frame_num, (uint32_t)(ptr-data), nal_size); - - switch (nal_unit_type) { - case HEVC_NAL_TYPE_VPS_NUT: - case HEVC_NAL_TYPE_SPS_NUT: - case HEVC_NAL_TYPE_PPS_NUT: - fwrite(ptr, 1, nal_size, of_prefix); - break; - case HEVC_NAL_TYPE_TRAIL_N: - case HEVC_NAL_TYPE_TRAIL_R: - case HEVC_NAL_TYPE_TSA_N: - case HEVC_NAL_TYPE_TSA_R: - case HEVC_NAL_TYPE_STSA_N: - case HEVC_NAL_TYPE_STSA_R: - case HEVC_NAL_TYPE_RADL_N: - case HEVC_NAL_TYPE_RADL_R: - case HEVC_NAL_TYPE_RASL_N: - case HEVC_NAL_TYPE_RASL_R: - case HEVC_NAL_TYPE_BLA_W_LP: - case HEVC_NAL_TYPE_BLA_W_RADL: - case HEVC_NAL_TYPE_BLA_N_LP: - case HEVC_NAL_TYPE_IDR_W_RADL: - case HEVC_NAL_TYPE_IDR_N_LP: - case HEVC_NAL_TYPE_CRA_NUT: { - // slice_segment_header - uint32_t first_slice_segment_in_pic_flag = bs_get(&bs, 1); - if (nal_unit_type >= HEVC_NAL_TYPE_BLA_W_LP && nal_unit_type <= HEVC_NAL_TYPE_RSV_IRAP_VCL23) { - uint32_t no_output_of_prior_pics_flag = bs_get(&bs, 1); - } - uint32_t slice_pic_parameter_set_id = bs_get(&bs, 1); - if (!first_slice_segment_in_pic_flag) { - // ... - break; - } - - if (!dependent_slice_segments_enabled_flag) { - for (int i=0; i 4); - - const uint8_t* data = (const uint8_t*)mmap(NULL, file_size, PROT_READ, MAP_PRIVATE, fd, 0); - assert(data != MAP_FAILED); - - if (strcmp(file_type, "hevc") == 0) { - hevc_index(data, file_size, of_prefix, of_index); - } else if (strcmp(file_type, "h264") == 0) { - h264_index(data, file_size, of_prefix, of_index); - } else { - assert(false); - } - - munmap((void*)data, file_size); - close(fd); - - return 0; -}