from __future__ import annotations import argparse import json import os import struct import sys from dataclasses import dataclass from typing import Dict, Iterator, List, Optional, Tuple METADATA_TYPE_ITUT_T35 = 4 HDR10PLUS_COUNTRY_CODE = 0xB5 HDR10PLUS_PROVIDER_CODE = 0x003C HDR10PLUS_ORIENTED_CODE = 0x0001 HDR10PLUS_APP_ID = 0x04 HDR10PLUS_APP_MODE = 0x01 OBU_METADATA = 5 DIST_INDEX_CLASSIC = [1, 5, 10, 25, 50, 75, 90, 95, 99] MATROSKA_CONTAINER_IDS = {0x1A45DFA3, 0x18538067, 0x1549A966, 0x1654AE6B, 0xAE, 0x1F43B675, 0xA0} MATROSKA_TRACKS_ID = 0x1654AE6B MATROSKA_TRACK_ENTRY_ID = 0xAE MATROSKA_TRACK_NUMBER_ID = 0xD7 MATROSKA_TRACK_TYPE_ID = 0x83 MATROSKA_CODEC_ID_ID = 0x86 MATROSKA_CLUSTER_ID = 0x1F43B675 MATROSKA_SIMPLE_BLOCK_ID = 0xA3 MATROSKA_BLOCK_GROUP_ID = 0xA0 MATROSKA_BLOCK_ID = 0xA1 def progress_bar(percent: float, width: int = 50) -> None: percent = max(0.0, min(100.0, float(percent))) filled = int(width * percent / 100.0) bar = "■" * filled + " " * (width - filled) sys.stderr.write(f"\r[{bar}] {percent:6.2f}%") sys.stderr.flush() def progress_done() -> None: progress_bar(100.0) sys.stderr.write("\n") sys.stderr.flush() def status(message: str) -> None: sys.stderr.write(message + "\n") sys.stderr.flush() def u16be(b: bytes) -> int: return (b[0] << 8) | b[1] def clamp(v: int, lo: int, hi: int) -> int: return lo if v < lo else hi if v > hi else v def leb128_read(data: bytes, off: int) -> Tuple[int, int]: val = 0 shift = 0 i = 0 while True: if off + i >= len(data): raise ValueError("LEB128 truncated") byte = data[off + i] val |= (byte & 0x7F) << shift i += 1 if (byte & 0x80) == 0: break shift += 7 if shift > 63: raise ValueError("LEB128 too long") return val, i class BitReader: def __init__(self, data: bytes): self.data = data self.bitpos = 0 def remaining_bits(self) -> int: return len(self.data) * 8 - self.bitpos def read_bits(self, n: int) -> int: if n <= 0: return 0 if self.remaining_bits() < n: raise ValueError("BitReader underflow") start = self.bitpos end = start + n byte_start = start >> 3 byte_end = (end + 7) >> 3 value = int.from_bytes(self.data[byte_start:byte_end], "big") extra_right_bits = (byte_end << 3) - end value >>= extra_right_bits value &= (1 << n) - 1 self.bitpos = end return value def read_u8(self) -> int: return self.read_bits(8) @dataclass class IvfHeader: width: int height: int timebase_num: int timebase_den: int frames: int @dataclass class Obu: obu_type: int header: bytes ext: bytes payload: bytes @dataclass class Box: type: bytes start: int size: int header_size: int end: int @dataclass class FragmentDefaults: base_data_offset: Optional[int] = None default_sample_duration: Optional[int] = None default_sample_size: Optional[int] = None default_sample_flags: Optional[int] = None @dataclass class SampleRef: offset: int size: int @dataclass class MatroskaElement: element_id: int data_start: int data_size: Optional[int] data_end: int header_start: int header_end: int @dataclass class MatroskaTrack: number: int track_type: int codec_id: str def read_ivf_header(fp) -> IvfHeader: header = fp.read(32) if len(header) != 32: raise ValueError("IVF header truncated") if header[0:4] != b"DKIF": raise ValueError("Not an IVF file") if header[8:12] != b"AV01": raise ValueError("IVF stream is not AV1") width = struct.unpack_from(" Iterator[Tuple[int, bytes, int]]: with open(path, "rb") as fp: header = read_ivf_header(fp) index = 0 while True: frame_header = fp.read(12) if not frame_header: return if len(frame_header) != 12: raise ValueError("IVF frame header truncated") size = struct.unpack_from(" List[Obu]: obus: List[Obu] = [] off = 0 n = len(frame_payload) while off < n: first = frame_payload[off] off += 1 forbidden = (first >> 7) & 1 if forbidden: raise ValueError("Invalid OBU header") obu_type = (first >> 3) & 0x0F ext_flag = (first >> 2) & 1 has_size_field = (first >> 1) & 1 ext = b"" if ext_flag: if off >= n: raise ValueError("OBU extension truncated") ext = frame_payload[off:off + 1] off += 1 if not has_size_field: raise ValueError("OBU has_size_field=0 is not supported") obu_size, used = leb128_read(frame_payload, off) off += used if off + obu_size > n: raise ValueError("OBU payload truncated") payload = frame_payload[off:off + obu_size] off += obu_size obus.append(Obu(obu_type=obu_type, header=bytes([first]), ext=ext, payload=payload)) return obus def is_hdr10plus_t35(t35: bytes) -> bool: if len(t35) < 7: return False if t35[0] != HDR10PLUS_COUNTRY_CODE: return False provider = u16be(t35[1:3]) oriented = u16be(t35[3:5]) app_id = t35[5] app_mode = t35[6] return provider == HDR10PLUS_PROVIDER_CODE and oriented == HDR10PLUS_ORIENTED_CODE and app_id == HDR10PLUS_APP_ID and app_mode == HDR10PLUS_APP_MODE def extract_t35_from_metadata_obu_payload(metadata_payload: bytes) -> Optional[bytes]: try: mtype, used = leb128_read(metadata_payload, 0) except Exception: return None if mtype != METADATA_TYPE_ITUT_T35: return None t35 = metadata_payload[used:] return t35 if is_hdr10plus_t35(t35) else None def extract_t35_from_av1_sample(sample_payload: bytes, sample_index: int) -> bytes: obus = parse_obus(sample_payload) for obu in obus: if obu.obu_type == OBU_METADATA: t35 = extract_t35_from_metadata_obu_payload(obu.payload) if t35 is not None: return t35 raise ValueError(f"No HDR10+ T.35 metadata found in frame/sample {sample_index}") def decode_hdr10plus_t35_av1(t35: bytes) -> Dict: if not is_hdr10plus_t35(t35): raise ValueError("Not HDR10+ T.35") br = BitReader(t35[7:]) num_windows = br.read_bits(2) targeted_raw = br.read_bits(27) tsd_actual_flag = br.read_bits(1) if tsd_actual_flag: rows = br.read_bits(5) cols = br.read_bits(5) for _r in range(rows): for _c in range(cols): _ = br.read_bits(4) max_scl_raw = [0, 0, 0] average_raw = 0 dist_map: Dict[int, int] = {} nw = num_windows if num_windows else 1 for _w in range(nw): max_scl_raw = [br.read_bits(17) for _ in range(3)] average_raw = br.read_bits(17) num_dists = br.read_bits(4) dist_map = {} for _ in range(num_dists): pct = br.read_bits(7) val = br.read_bits(17) dist_map[int(pct)] = int(val) _ = br.read_bits(10) mastering_flag = br.read_bits(1) if mastering_flag: rows = br.read_bits(5) cols = br.read_bits(5) for _r in range(rows): for _c in range(cols): _ = br.read_bits(4) knee_x = 0 knee_y = 0 anchors: List[int] = [0] * 9 for _w in range(nw): tone_mapping_flag = br.read_bits(1) if tone_mapping_flag: knee_x = br.read_bits(12) knee_y = br.read_bits(12) n_anchors = br.read_bits(4) raw_anchors = [] for _ in range(n_anchors): raw_anchors.append(br.read_bits(10)) anchors = (raw_anchors[:9] + [0] * 9)[:9] color_flag = br.read_bits(1) if br.remaining_bits() >= 1 else 0 if color_flag and br.remaining_bits() >= 6: _ = br.read_bits(6) dist_vals = [int(dist_map.get(i, 0)) for i in DIST_INDEX_CLASSIC] targeted_out = int(targeted_raw // 64) if targeted_raw > 10000 else int(targeted_raw) max_scl_out = [int(x // 32) for x in max_scl_raw] if any(x > 100000 for x in max_scl_raw) else [int(x) for x in max_scl_raw] average_out = int(average_raw // 32) if average_raw > 100000 else int(average_raw) return { "NumberOfWindows": int(nw), "TargetedSystemDisplayMaximumLuminance": targeted_out, "AverageRGB": average_out, "MaxScl": max_scl_out, "DistributionIndex": list(DIST_INDEX_CLASSIC), "DistributionValues": dist_vals, "BezierCurveData": { "Anchors": [int(x) for x in anchors], "KneePointX": int(knee_x), "KneePointY": int(knee_y), }, } def metadata_signature(decoded: Dict) -> Tuple: return ( decoded.get("NumberOfWindows"), decoded.get("TargetedSystemDisplayMaximumLuminance"), decoded.get("AverageRGB"), tuple(decoded.get("MaxScl", [])), tuple(decoded.get("DistributionValues", [])), decoded.get("BezierCurveData", {}).get("KneePointX"), decoded.get("BezierCurveData", {}).get("KneePointY"), tuple(decoded.get("BezierCurveData", {}).get("Anchors", [])), ) def build_classic_json_from_t35_list(t35_by_frame: List[bytes]) -> Dict: scene_info: List[Dict] = [] scene_id = 0 scene_frame_index = 0 prev_sig: Optional[Tuple] = None for i, t35 in enumerate(t35_by_frame): decoded = decode_hdr10plus_t35_av1(t35) sig = metadata_signature(decoded) if prev_sig is None: scene_id = 0 scene_frame_index = 0 else: if sig != prev_sig: scene_id += 1 scene_frame_index = 0 else: scene_frame_index += 1 entry = { "BezierCurveData": decoded["BezierCurveData"], "LuminanceParameters": { "AverageRGB": decoded["AverageRGB"], "LuminanceDistributions": { "DistributionIndex": decoded["DistributionIndex"], "DistributionValues": decoded["DistributionValues"], }, "MaxScl": decoded["MaxScl"], }, "NumberOfWindows": decoded["NumberOfWindows"], "TargetedSystemDisplayMaximumLuminance": decoded["TargetedSystemDisplayMaximumLuminance"], "SceneFrameIndex": int(scene_frame_index), "SceneId": int(scene_id), "SequenceFrameIndex": int(i), } scene_info.append(entry) prev_sig = sig scene_first: List[int] = [] scene_len: List[int] = [] if scene_info: cur = scene_info[0]["SceneId"] start = 0 for idx in range(1, len(scene_info)): if scene_info[idx]["SceneId"] != cur: scene_first.append(start) scene_len.append(idx - start) start = idx cur = scene_info[idx]["SceneId"] scene_first.append(start) scene_len.append(len(scene_info) - start) return { "JSONInfo": {"HDR10plusProfile": "B", "Version": "1.0"}, "SceneInfo": scene_info, "SceneInfoSummary": { "SceneFirstFrameIndex": scene_first, "SceneFrameNumbers": scene_len, }, "ToolInfo": {"Tool": "pyav1plus", "Version": "1.1.0"}, } def read_box_header(fp, file_size: int) -> Optional[Box]: start = fp.tell() if start >= file_size: return None header = fp.read(8) if not header: return None if len(header) != 8: raise ValueError(f"MP4 box header truncated at offset {start}") size32, box_type = struct.unpack(">I4s", header) header_size = 8 if size32 == 1: ext = fp.read(8) if len(ext) != 8: raise ValueError(f"MP4 extended box size truncated at offset {start}") size = struct.unpack(">Q", ext)[0] header_size = 16 elif size32 == 0: size = file_size - start else: size = size32 if size < header_size: raise ValueError(f"Invalid MP4 box size at offset {start}") return Box(box_type, start, size, header_size, start + size) def iter_top_level_boxes(fp, file_size: int) -> Iterator[Box]: fp.seek(0) while fp.tell() < file_size: box = read_box_header(fp, file_size) if box is None: break yield box fp.seek(box.end) def iter_child_boxes(fp, start: int, end: int, file_size: int) -> Iterator[Box]: fp.seek(start) while fp.tell() < end: box = read_box_header(fp, file_size) if box is None: break if box.end > end: raise ValueError(f"Child MP4 box {box.type!r} exceeds parent bounds") yield box fp.seek(box.end) def parse_tfhd(data: bytes, moof_start: int) -> FragmentDefaults: if len(data) < 8: raise ValueError("tfhd box is truncated") version_flags = struct.unpack_from(">I", data, 0)[0] flags = version_flags & 0xFFFFFF offset = 8 defaults = FragmentDefaults() if flags & 0x000001: defaults.base_data_offset = struct.unpack_from(">Q", data, offset)[0] offset += 8 if flags & 0x000002: offset += 4 if flags & 0x000008: defaults.default_sample_duration = struct.unpack_from(">I", data, offset)[0] offset += 4 if flags & 0x000010: defaults.default_sample_size = struct.unpack_from(">I", data, offset)[0] offset += 4 if flags & 0x000020: defaults.default_sample_flags = struct.unpack_from(">I", data, offset)[0] offset += 4 if flags & 0x020000: defaults.base_data_offset = moof_start return defaults def parse_trun(data: bytes, defaults: FragmentDefaults, moof_start: int) -> List[SampleRef]: if len(data) < 8: raise ValueError("trun box is truncated") version_flags = struct.unpack_from(">I", data, 0)[0] flags = version_flags & 0xFFFFFF sample_count = struct.unpack_from(">I", data, 4)[0] offset = 8 data_offset = 0 if flags & 0x000001: data_offset = struct.unpack_from(">i", data, offset)[0] offset += 4 if flags & 0x000004: offset += 4 current = (defaults.base_data_offset if defaults.base_data_offset is not None else moof_start) + data_offset samples: List[SampleRef] = [] for _ in range(sample_count): if flags & 0x000100: offset += 4 if flags & 0x000200: sample_size = struct.unpack_from(">I", data, offset)[0] offset += 4 else: sample_size = defaults.default_sample_size if sample_size is None: raise ValueError("trun does not contain sample sizes and tfhd has no default_sample_size") if flags & 0x000400: offset += 4 if flags & 0x000800: offset += 4 samples.append(SampleRef(int(current), int(sample_size))) current += int(sample_size) return samples def collect_fragmented_mp4_samples(path: str) -> List[SampleRef]: samples: List[SampleRef] = [] file_size = os.path.getsize(path) with open(path, "rb") as fp: for top in iter_top_level_boxes(fp, file_size): if top.type != b"moof": continue for moof_child in iter_child_boxes(fp, top.start + top.header_size, top.end, file_size): if moof_child.type != b"traf": continue defaults = FragmentDefaults() trun_payloads: List[bytes] = [] for traf_child in iter_child_boxes(fp, moof_child.start + moof_child.header_size, moof_child.end, file_size): fp.seek(traf_child.start + traf_child.header_size) payload = fp.read(traf_child.size - traf_child.header_size) if traf_child.type == b"tfhd": defaults = parse_tfhd(payload, top.start) elif traf_child.type == b"trun": trun_payloads.append(payload) for payload in trun_payloads: samples.extend(parse_trun(payload, defaults, top.start)) return samples def parse_stsz(data: bytes) -> List[int]: sample_size = struct.unpack_from(">I", data, 4)[0] sample_count = struct.unpack_from(">I", data, 8)[0] if sample_size: return [int(sample_size)] * sample_count return [struct.unpack_from(">I", data, 12 + i * 4)[0] for i in range(sample_count)] def parse_stco(data: bytes) -> List[int]: count = struct.unpack_from(">I", data, 4)[0] return [struct.unpack_from(">I", data, 8 + i * 4)[0] for i in range(count)] def parse_co64(data: bytes) -> List[int]: count = struct.unpack_from(">I", data, 4)[0] return [struct.unpack_from(">Q", data, 8 + i * 8)[0] for i in range(count)] def parse_stsc(data: bytes) -> List[Tuple[int, int, int]]: count = struct.unpack_from(">I", data, 4)[0] entries = [] for i in range(count): entries.append(struct.unpack_from(">III", data, 8 + i * 12)) return entries def collect_unfragmented_mp4_samples(path: str) -> List[SampleRef]: file_size = os.path.getsize(path) container_boxes = {b"moov", b"trak", b"mdia", b"minf", b"stbl", b"edts", b"mvex", b"meta"} def walk(fp, start: int, end: int): fp.seek(start) while fp.tell() < end: box = read_box_header(fp, file_size) if box is None: return None if box.type == b"stbl": stsz = stco = co64 = stsc = None for child in iter_child_boxes(fp, box.start + box.header_size, box.end, file_size): fp.seek(child.start + child.header_size) payload = fp.read(child.size - child.header_size) if child.type == b"stsz": stsz = parse_stsz(payload) elif child.type == b"stco": stco = parse_stco(payload) elif child.type == b"co64": co64 = parse_co64(payload) elif child.type == b"stsc": stsc = parse_stsc(payload) chunk_offsets = co64 if co64 is not None else stco if stsz and chunk_offsets and stsc: return stsz, chunk_offsets, stsc elif box.type in container_boxes: result = walk(fp, box.start + box.header_size, box.end) if result is not None: return result fp.seek(box.end) return None with open(path, "rb") as fp: result = walk(fp, 0, file_size) if result is None: return [] sizes, chunk_offsets, stsc_entries = result samples: List[SampleRef] = [] sample_index = 0 for chunk_index, chunk_offset in enumerate(chunk_offsets, start=1): active = stsc_entries[0] for entry_index, entry in enumerate(stsc_entries): next_first = stsc_entries[entry_index + 1][0] if entry_index + 1 < len(stsc_entries) else 1 << 60 if entry[0] <= chunk_index < next_first: active = entry break current = int(chunk_offset) for _ in range(active[1]): if sample_index >= len(sizes): break size = int(sizes[sample_index]) samples.append(SampleRef(current, size)) current += size sample_index += 1 return samples def collect_mp4_samples(path: str) -> List[SampleRef]: samples = collect_fragmented_mp4_samples(path) if samples: return samples samples = collect_unfragmented_mp4_samples(path) if samples: return samples raise ValueError("No MP4 AV1 samples were found") def iter_mp4_samples(path: str) -> Iterator[Tuple[int, bytes, int]]: samples = collect_mp4_samples(path) total = len(samples) with open(path, "rb") as fp: for index, sample in enumerate(samples): fp.seek(sample.offset) payload = fp.read(sample.size) if len(payload) != sample.size: raise ValueError(f"MP4 sample {index} is truncated") yield index, payload, total def read_ebml_id_from_file(fp) -> Optional[Tuple[int, int]]: first_raw = fp.read(1) if not first_raw: return None first = first_raw[0] mask = 0x80 length = 1 while length <= 4 and not (first & mask): mask >>= 1 length += 1 if length > 4: raise ValueError("Invalid EBML element ID") rest = fp.read(length - 1) if len(rest) != length - 1: raise ValueError("Truncated EBML element ID") return int.from_bytes(first_raw + rest, "big"), length def read_ebml_size_from_file(fp) -> Tuple[Optional[int], int]: first_raw = fp.read(1) if not first_raw: raise ValueError("Truncated EBML element size") first = first_raw[0] mask = 0x80 length = 1 while length <= 8 and not (first & mask): mask >>= 1 length += 1 if length > 8: raise ValueError("Invalid EBML element size") value = first & (mask - 1) rest = fp.read(length - 1) if len(rest) != length - 1: raise ValueError("Truncated EBML element size") for b in rest: value = (value << 8) | b unknown = value == ((1 << (7 * length)) - 1) return None if unknown else value, length def read_matroska_element(fp, limit: int) -> Optional[MatroskaElement]: header_start = fp.tell() if header_start >= limit: return None id_result = read_ebml_id_from_file(fp) if id_result is None: return None element_id, id_len = id_result data_size, size_len = read_ebml_size_from_file(fp) data_start = fp.tell() data_end = limit if data_size is None else data_start + data_size if data_end > limit: data_end = limit return MatroskaElement(element_id, data_start, data_size, data_end, header_start, data_start) def read_unsigned_element(fp, element: MatroskaElement) -> int: fp.seek(element.data_start) data = fp.read(element.data_end - element.data_start) if not data: return 0 return int.from_bytes(data, "big") def read_string_element(fp, element: MatroskaElement) -> str: fp.seek(element.data_start) data = fp.read(element.data_end - element.data_start) return data.decode("utf-8", "replace") def parse_matroska_track_entry(fp, entry: MatroskaElement) -> Optional[MatroskaTrack]: number = 0 track_type = 0 codec_id = "" fp.seek(entry.data_start) while fp.tell() < entry.data_end: child = read_matroska_element(fp, entry.data_end) if child is None: break if child.element_id == MATROSKA_TRACK_NUMBER_ID: number = read_unsigned_element(fp, child) elif child.element_id == MATROSKA_TRACK_TYPE_ID: track_type = read_unsigned_element(fp, child) elif child.element_id == MATROSKA_CODEC_ID_ID: codec_id = read_string_element(fp, child) fp.seek(child.data_end) if number <= 0: return None return MatroskaTrack(number=number, track_type=track_type, codec_id=codec_id) def parse_matroska_tracks(fp, tracks_element: MatroskaElement) -> Dict[int, MatroskaTrack]: tracks: Dict[int, MatroskaTrack] = {} fp.seek(tracks_element.data_start) while fp.tell() < tracks_element.data_end: child = read_matroska_element(fp, tracks_element.data_end) if child is None: break if child.element_id == MATROSKA_TRACK_ENTRY_ID: track = parse_matroska_track_entry(fp, child) if track is not None: tracks[track.number] = track fp.seek(child.data_end) return tracks def collect_matroska_tracks(path: str) -> Dict[int, MatroskaTrack]: file_size = os.path.getsize(path) tracks: Dict[int, MatroskaTrack] = {} def walk(fp, start: int, end: int) -> None: fp.seek(start) while fp.tell() < end: element = read_matroska_element(fp, end) if element is None: return if element.element_id == MATROSKA_TRACKS_ID: tracks.update(parse_matroska_tracks(fp, element)) return if element.element_id in MATROSKA_CONTAINER_IDS and element.element_id != MATROSKA_CLUSTER_ID: walk(fp, element.data_start, element.data_end) if tracks: return fp.seek(element.data_end) with open(path, "rb") as fp: walk(fp, 0, file_size) return tracks def read_block_track_number(data: bytes, off: int = 0) -> Tuple[int, int]: if off >= len(data): raise ValueError("Matroska block is truncated") first = data[off] mask = 0x80 length = 1 while length <= 8 and not (first & mask): mask >>= 1 length += 1 if length > 8 or off + length > len(data): raise ValueError("Invalid Matroska block track number") value = first & (mask - 1) for b in data[off + 1:off + length]: value = (value << 8) | b return value, length def split_xiph_lacing(data: bytes, off: int, frame_count: int) -> List[bytes]: sizes: List[int] = [] cursor = off for _ in range(frame_count - 1): size = 0 while True: if cursor >= len(data): raise ValueError("Xiph lacing is truncated") b = data[cursor] cursor += 1 size += b if b != 255: break sizes.append(size) used = sum(sizes) remaining = len(data) - cursor - used if remaining < 0: raise ValueError("Invalid Xiph lacing sizes") sizes.append(remaining) frames = [] for size in sizes: frames.append(data[cursor:cursor + size]) cursor += size return frames def read_ebml_vint_from_bytes(data: bytes, off: int) -> Tuple[int, int]: if off >= len(data): raise ValueError("EBML lacing size is truncated") first = data[off] mask = 0x80 length = 1 while length <= 8 and not (first & mask): mask >>= 1 length += 1 if length > 8 or off + length > len(data): raise ValueError("Invalid EBML lacing size") value = first & (mask - 1) for b in data[off + 1:off + length]: value = (value << 8) | b return value, length def signed_ebml_value(value: int, length: int) -> int: bias = (1 << (7 * length - 1)) - 1 return value - bias def split_ebml_lacing(data: bytes, off: int, frame_count: int) -> List[bytes]: sizes: List[int] = [] cursor = off first_size, used = read_ebml_vint_from_bytes(data, cursor) cursor += used sizes.append(first_size) previous = first_size for _ in range(frame_count - 2): raw, used = read_ebml_vint_from_bytes(data, cursor) diff = signed_ebml_value(raw, used) cursor += used current = previous + diff if current < 0: raise ValueError("Invalid EBML lacing size") sizes.append(current) previous = current used_total = sum(sizes) remaining = len(data) - cursor - used_total if remaining < 0: raise ValueError("Invalid EBML lacing sizes") sizes.append(remaining) frames = [] for size in sizes: frames.append(data[cursor:cursor + size]) cursor += size return frames def split_fixed_lacing(data: bytes, off: int, frame_count: int) -> List[bytes]: remaining = len(data) - off if frame_count <= 0 or remaining % frame_count: raise ValueError("Invalid fixed lacing size") size = remaining // frame_count return [data[off + i * size:off + (i + 1) * size] for i in range(frame_count)] def parse_matroska_block_frames(data: bytes) -> Tuple[int, List[bytes]]: track_number, used = read_block_track_number(data, 0) cursor = used + 2 if cursor >= len(data): raise ValueError("Matroska block is truncated") flags = data[cursor] cursor += 1 lacing = (flags >> 1) & 0x03 if lacing == 0: return track_number, [data[cursor:]] if cursor >= len(data): raise ValueError("Matroska laced block is truncated") frame_count = data[cursor] + 1 cursor += 1 if lacing == 1: return track_number, split_xiph_lacing(data, cursor, frame_count) if lacing == 2: return track_number, split_fixed_lacing(data, cursor, frame_count) return track_number, split_ebml_lacing(data, cursor, frame_count) def iter_matroska_block_payloads(path: str) -> Iterator[Tuple[int, bytes, int]]: file_size = os.path.getsize(path) tracks = collect_matroska_tracks(path) av1_tracks = {number for number, track in tracks.items() if track.track_type == 1 and track.codec_id.upper() == "V_AV1"} if not av1_tracks and tracks: av1_tracks = {number for number, track in tracks.items() if track.codec_id.upper() == "V_AV1"} index = 0 with open(path, "rb") as fp: stack: List[Tuple[int, bool]] = [(file_size, False)] while stack: limit, in_cluster = stack.pop() while fp.tell() < limit: element = read_matroska_element(fp, limit) if element is None: break if element.element_id == MATROSKA_CLUSTER_ID: stack.append((limit, in_cluster)) fp.seek(element.data_start) stack.append((element.data_end, True)) break if in_cluster and element.element_id in {MATROSKA_SIMPLE_BLOCK_ID, MATROSKA_BLOCK_ID}: fp.seek(element.data_start) block_data = fp.read(element.data_end - element.data_start) track_number, frames = parse_matroska_block_frames(block_data) if not av1_tracks or track_number in av1_tracks: for frame in frames: yield index, frame, 0 index += 1 elif element.element_id in MATROSKA_CONTAINER_IDS and element.element_id != MATROSKA_BLOCK_GROUP_ID: stack.append((limit, in_cluster)) fp.seek(element.data_start) stack.append((element.data_end, in_cluster)) break elif in_cluster and element.element_id == MATROSKA_BLOCK_GROUP_ID: stack.append((limit, in_cluster)) fp.seek(element.data_start) stack.append((element.data_end, True)) break fp.seek(element.data_end) def iter_mkv_samples(path: str) -> Iterator[Tuple[int, bytes, int]]: yield from iter_matroska_block_payloads(path) def detect_input_format(path: str) -> str: with open(path, "rb") as fp: head = fp.read(16) if head.startswith(b"DKIF"): return "ivf" if head.startswith(bytes.fromhex("1a45dfa3")): return "mkv" if len(head) >= 8 and head[4:8] in {b"ftyp", b"moov", b"moof", b"free", b"wide", b"mdat"}: return "mp4" raise ValueError("Unsupported input format. Expected AV1 IVF, AV1 MP4/fMP4, or AV1 MKV.") def output_path_for_input(path: str) -> str: base, _ext = os.path.splitext(path) return base + ".json" def build_iterator(path: str, input_format: str) -> Iterator[Tuple[int, bytes, int]]: if input_format == "ivf": return iter_ivf_samples(path) if input_format == "mp4": return iter_mp4_samples(path) if input_format == "mkv": return iter_mkv_samples(path) raise ValueError("Unsupported input format") def extract(path: str, output: str) -> None: input_format = detect_input_format(path) iterator = build_iterator(path, input_format) t35_by_frame: List[bytes] = [] scanned = 0 total_hint = os.path.getsize(path) if input_format == "mkv" else 0 status(f"Input format: {input_format.upper()}") status("Scanning AV1 samples for HDR10+ metadata...") last_progress = -1.0 last_file_pos = 0 for index, payload, total in iterator: scanned += 1 t35_by_frame.append(extract_t35_from_av1_sample(payload, index)) if total > 0: current_progress = scanned * 100.0 / total elif input_format == "mkv" and total_hint > 0: current_progress = min(99.0, last_file_pos * 100.0 / total_hint) else: current_progress = min(99.0, (scanned % 1000) / 10.0) if current_progress >= 100.0 or current_progress - last_progress >= 0.10: progress_bar(current_progress) last_progress = current_progress last_file_pos = min(total_hint, last_file_pos + len(payload)) if scanned == 0: raise ValueError("No AV1 samples were found") progress_done() status("Reordering metadata... Done.") status("Reading parsed dynamic metadata... Done.") classic = build_classic_json_from_t35_list(t35_by_frame) status("Generating and writing metadata to JSON file...") with open(output, "w", encoding="utf-8") as fp: json.dump(classic, fp, indent=2) status("Generating and writing metadata to JSON file... Done.") status(f"Frames written: {len(t35_by_frame)}") status(f"Output: {output}") def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(prog="pyav1plus.py", description="Extract HDR10+ metadata from AV1 IVF, MP4/fMP4, or MKV.") parser.add_argument("input", help="Input .ivf, .mp4, .m4v, .mkv, or .webm file") parser.add_argument("-o", "--output", help="Output JSON path. Default: input basename with .json") return parser def main() -> None: args = build_parser().parse_args() output = args.output or output_path_for_input(args.input) try: extract(args.input, output) except KeyboardInterrupt: sys.stderr.write("\nInterrupted.\n") raise SystemExit(130) except Exception as exc: sys.stderr.write(f"\nERROR: {exc}\n") raise SystemExit(1) if __name__ == "__main__": main()