mirror of
https://github.com/Hugoved/pyav1plus.git
synced 2026-07-15 18:10:05 +02:00
Add files via upload
This commit is contained in:
Hugoved
+959
@@ -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("<H", header, 12)[0]
|
||||
height = struct.unpack_from("<H", header, 14)[0]
|
||||
tb_den = struct.unpack_from("<I", header, 16)[0]
|
||||
tb_num = struct.unpack_from("<I", header, 20)[0]
|
||||
frames = struct.unpack_from("<I", header, 24)[0]
|
||||
return IvfHeader(width, height, tb_num, tb_den, frames)
|
||||
|
||||
def iter_ivf_samples(path: str) -> 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("<I", frame_header, 0)[0]
|
||||
payload = fp.read(size)
|
||||
if len(payload) != size:
|
||||
raise ValueError("IVF frame payload truncated")
|
||||
yield index, payload, header.frames
|
||||
index += 1
|
||||
|
||||
def parse_obus(frame_payload: bytes) -> 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()
|
||||
Reference in New Issue
Block a user