From 7260b1c3d6dc3aaaa9fb6a92baebc2cf389e2636 Mon Sep 17 00:00:00 2001 From: Hugoved Date: Mon, 11 May 2026 09:52:13 +0200 Subject: [PATCH] Add files via upload --- pyav1plus.py | 959 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 959 insertions(+) create mode 100644 pyav1plus.py diff --git a/pyav1plus.py b/pyav1plus.py new file mode 100644 index 0000000..59d4b4f --- /dev/null +++ b/pyav1plus.py @@ -0,0 +1,959 @@ +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() \ No newline at end of file