189 lines
8.7 KiB
Python
189 lines
8.7 KiB
Python
"""VIDEO-003: сравнение H.264 через PyAV/libx264 на исходнике VIDEO-001."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
from fractions import Fraction
|
||
import hashlib
|
||
import json
|
||
import math
|
||
from pathlib import Path
|
||
import statistics
|
||
|
||
import av
|
||
import numpy as np
|
||
|
||
from tools.video_compare import recovery_after_drop, source_frames
|
||
|
||
|
||
ROOT = Path(__file__).resolve().parents[1]
|
||
SOURCE = ROOT / "data/raw/video001/udacity_project_video.mp4"
|
||
BASELINE = ROOT / "data/processed/video001/comparison_traffic_320k.json"
|
||
DEFAULT_OUTPUT = ROOT / "data/processed/video003/codec_results.json"
|
||
VIDEO_MODES = (("NORMAL", 640, 360, 10), ("HIGH_QUALITY", 854, 480, 10),
|
||
("720P5", 1280, 720, 5))
|
||
PROFILES = ("CURRENT", "CURRENT_VBV", "SHORT_GOP", "INTRA_REFRESH")
|
||
BITRATES = (300_000, 400_000, 500_000)
|
||
|
||
|
||
def _encoder(width: int, height: int, fps: int, profile: str,
|
||
target_bps: int | None) -> tuple[av.CodecContext, dict]:
|
||
if profile not in PROFILES:
|
||
raise ValueError("неизвестный профиль H.264")
|
||
if profile == "CURRENT" and target_bps is not None:
|
||
raise ValueError("текущий профиль использует только исходный VBV 320k")
|
||
if profile != "CURRENT" and target_bps not in BITRATES:
|
||
raise ValueError("для профиля VBV нужен target 300/400/500 кбит/с")
|
||
encoder = av.CodecContext.create("libx264", "w")
|
||
encoder.width, encoder.height = width, height
|
||
encoder.pix_fmt = "yuv420p"
|
||
encoder.time_base = Fraction(1, fps)
|
||
encoder.framerate = Fraction(fps, 1)
|
||
encoder.thread_count = 1
|
||
encoder.max_b_frames = 0
|
||
gop_size = math.ceil(fps / 2) if profile == "SHORT_GOP" else fps
|
||
encoder.gop_size = gop_size
|
||
if profile != "CURRENT":
|
||
encoder.bit_rate = target_bps
|
||
maxrate_bps = 320_000 if profile == "CURRENT" else target_bps
|
||
buffer_bps = maxrate_bps if profile == "CURRENT" else round(target_bps * 0.2)
|
||
intra_refresh = profile == "INTRA_REFRESH"
|
||
parameters = (
|
||
f"keyint={gop_size}:min-keyint={gop_size}:scenecut=0:open-gop=0:"
|
||
f"bframes=0:rc-lookahead=0:sync-lookahead=0:repeat-headers=1:"
|
||
f"vbv-maxrate={maxrate_bps // 1000}:vbv-bufsize={buffer_bps // 1000}:"
|
||
f"intra-refresh={int(intra_refresh)}"
|
||
)
|
||
options = {"preset": "veryfast", "tune": "zerolatency", "crf": "28",
|
||
"x264-params": parameters}
|
||
encoder.options = options.copy()
|
||
encoder.open()
|
||
settings = {
|
||
"encoder_api": f"PyAV {av.__version__}",
|
||
"ffmpeg_libraries": {name: ".".join(map(str, version))
|
||
for name, version in av.library_versions.items()},
|
||
"ffmpeg_cli_used": False,
|
||
"codec_name": "libx264", "resolution": [width, height], "frame_rate": fps,
|
||
"pixel_format": "yuv420p", "bitrate_target_bps": target_bps,
|
||
"max_bitrate_bps": maxrate_bps, "buffer_size_bits": buffer_bps,
|
||
"buffer_duration_s": buffer_bps / maxrate_bps,
|
||
"gop_size_frames": gop_size, "b_frame_count": 0,
|
||
"preset": options["preset"], "tune": options["tune"],
|
||
"profile_requested": None, "profile_actual": encoder.profile,
|
||
"intra_refresh_option": intra_refresh,
|
||
"time_base": f"1/{fps}", "codec_options": options,
|
||
"bitrate_property_after_open_bps": encoder.bit_rate,
|
||
}
|
||
return encoder, settings
|
||
|
||
|
||
def _peak_rate(frame_sizes: list[int], fps: int, window_ms: int) -> float:
|
||
"""Максимум битов по полуоткрытому временному окну длиной window_ms."""
|
||
frame_times = [index * 1000 / fps for index in range(len(frame_sizes))]
|
||
best_bits = 0
|
||
for start in frame_times:
|
||
bits = sum(size * 8 for time_ms, size in zip(frame_times, frame_sizes)
|
||
if start <= time_ms < start + window_ms - 1e-9)
|
||
best_bits = max(best_bits, bits)
|
||
return best_bits * 1000 / window_ms
|
||
|
||
|
||
def encode_case(frames: list[np.ndarray], fps: int, profile: str,
|
||
target_bps: int | None) -> tuple[dict, list[bytes]]:
|
||
height, width = frames[0].shape[:2]
|
||
encoder, settings = _encoder(width, height, fps, profile, target_bps)
|
||
packets: list[bytes] = []
|
||
keyframes: list[int] = []
|
||
for index, pixels in enumerate(frames):
|
||
frame = av.VideoFrame.from_ndarray(pixels, format="rgb24").reformat(format="yuv420p")
|
||
frame.pts = index
|
||
encoded = encoder.encode(frame)
|
||
if len(encoded) != 1:
|
||
raise ValueError("ожидался один немедленный H.264 packet на кадр")
|
||
if encoded[0].is_keyframe:
|
||
keyframes.append(index)
|
||
packets.append(bytes(encoded[0]))
|
||
if encoder.encode(None):
|
||
raise ValueError("кодировщик задержал пакеты вопреки low-latency режиму")
|
||
decoder = av.CodecContext.create("h264", "r")
|
||
decoder.open()
|
||
clean_hashes = {}
|
||
psnr_values = []
|
||
decoded_count = 0
|
||
for index, data in enumerate(packets):
|
||
packet = av.Packet(data)
|
||
packet.pts = packet.dts = index
|
||
packet.time_base = Fraction(1, fps)
|
||
decoded = decoder.decode(packet)
|
||
if len(decoded) != 1 or decoded[0].pts != index:
|
||
raise ValueError("чистое декодирование не дало кадр с ожидаемым номером")
|
||
pixels = decoded[0].to_ndarray(format="rgb24")
|
||
clean_hashes[index] = hashlib.sha256(pixels.tobytes()).hexdigest()
|
||
if index == 0:
|
||
settings["profile_actual"] = decoder.profile
|
||
error = float(np.mean((frames[index].astype(np.float32)
|
||
- pixels.astype(np.float32)) ** 2))
|
||
psnr_values.append(10 * math.log10(255 * 255 / max(error, 1e-12)))
|
||
decoded_count += 1
|
||
sizes = [len(packet) for packet in packets]
|
||
dropped_frame = min(fps + 1, len(packets) - 2)
|
||
recovery = recovery_after_drop(packets, dropped_frame, fps, clean_hashes)
|
||
recovered_at = recovery["first_exact_frame"]
|
||
result = {
|
||
"settings": settings,
|
||
"frame_count": len(packets), "decoded_frame_count": decoded_count,
|
||
"frame_sizes_bytes": sizes, "keyframe_positions": keyframes,
|
||
"actual_average_bitrate_bps": sum(sizes) * 8 * fps / len(sizes),
|
||
"peak_100ms_bitrate_bps": _peak_rate(sizes, fps, 100),
|
||
"peak_1s_bitrate_bps": _peak_rate(sizes, fps, 1000),
|
||
"max_frame_size_bytes": max(sizes),
|
||
"p95_frame_size_bytes": float(np.percentile(sizes, 95)),
|
||
"quality": {"mean_psnr_db": statistics.mean(psnr_values),
|
||
"min_psnr_db": min(psnr_values),
|
||
"operator_visual_review": "NOT_RUN"},
|
||
"loss_recovery": {"dropped_frame": dropped_frame,
|
||
"first_exact_frame": recovered_at,
|
||
"recovery_time_ms": ((recovered_at - dropped_frame) * 1000 / fps
|
||
if recovered_at is not None else None),
|
||
"decoder_errors": recovery["decoder_errors"]},
|
||
}
|
||
return result, packets
|
||
|
||
|
||
def run() -> dict:
|
||
baseline = json.loads(BASELINE.read_text(encoding="utf-8"))
|
||
digest = hashlib.sha256(SOURCE.read_bytes()).hexdigest()
|
||
if digest != baseline["source_sha256"]:
|
||
raise ValueError("исходное видео не совпадает с VIDEO-001")
|
||
results = []
|
||
for mode, width, height, fps in VIDEO_MODES:
|
||
for window in baseline["windows"]:
|
||
frames = source_frames(SOURCE, width, height, fps,
|
||
window["start_seconds"], window["duration_seconds"])
|
||
for profile in PROFILES:
|
||
targets = (None,) if profile == "CURRENT" else BITRATES
|
||
for target in targets:
|
||
item, _ = encode_case(frames, fps, profile, target)
|
||
item.update({"mode": mode, "window": window["name"],
|
||
"profile": profile, "target_bps": target})
|
||
results.append(item)
|
||
print(f"{mode}/{window['name']}/{profile}/{target}: "
|
||
f"{item['actual_average_bitrate_bps'] / 1000:.1f} кбит/с",
|
||
flush=True)
|
||
return {"source": str(SOURCE.relative_to(ROOT)), "source_sha256": digest,
|
||
"cases": results, "rf_tx": 0}
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(description=__doc__)
|
||
parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
|
||
args = parser.parse_args()
|
||
result = run()
|
||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||
args.output.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
print(f"VIDEO-003: {len(result['cases'])} H.264 прогонов; {args.output}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|