Initial Commit
This commit is contained in:
Andy
451
unshackle/core/tracks/video.py
Normal file
451
unshackle/core/tracks/video.py
Normal file
@@ -0,0 +1,451 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import math
|
||||
import re
|
||||
import subprocess
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
from langcodes import Language
|
||||
|
||||
from unshackle.core import binaries
|
||||
from unshackle.core.config import config
|
||||
from unshackle.core.tracks.subtitle import Subtitle
|
||||
from unshackle.core.tracks.track import Track
|
||||
from unshackle.core.utilities import FPS, get_boxes
|
||||
|
||||
|
||||
class Video(Track):
|
||||
class Codec(str, Enum):
|
||||
AVC = "H.264"
|
||||
HEVC = "H.265"
|
||||
VC1 = "VC-1"
|
||||
VP8 = "VP8"
|
||||
VP9 = "VP9"
|
||||
AV1 = "AV1"
|
||||
|
||||
@property
|
||||
def extension(self) -> str:
|
||||
return self.value.lower().replace(".", "").replace("-", "")
|
||||
|
||||
@staticmethod
|
||||
def from_mime(mime: str) -> Video.Codec:
|
||||
mime = mime.lower().strip().split(".")[0]
|
||||
if mime in (
|
||||
"avc1",
|
||||
"avc2",
|
||||
"avc3",
|
||||
"dva1",
|
||||
"dvav", # Dolby Vision
|
||||
):
|
||||
return Video.Codec.AVC
|
||||
if mime in (
|
||||
"hev1",
|
||||
"hev2",
|
||||
"hev3",
|
||||
"hvc1",
|
||||
"hvc2",
|
||||
"hvc3",
|
||||
"dvh1",
|
||||
"dvhe", # Dolby Vision
|
||||
"lhv1",
|
||||
"lhe1", # Layered
|
||||
):
|
||||
return Video.Codec.HEVC
|
||||
if mime == "vc-1":
|
||||
return Video.Codec.VC1
|
||||
if mime in ("vp08", "vp8"):
|
||||
return Video.Codec.VP8
|
||||
if mime in ("vp09", "vp9"):
|
||||
return Video.Codec.VP9
|
||||
if mime == "av01":
|
||||
return Video.Codec.AV1
|
||||
raise ValueError(f"The MIME '{mime}' is not a supported Video Codec")
|
||||
|
||||
@staticmethod
|
||||
def from_codecs(codecs: str) -> Video.Codec:
|
||||
for codec in codecs.lower().split(","):
|
||||
codec = codec.strip()
|
||||
mime = codec.split(".")[0]
|
||||
try:
|
||||
return Video.Codec.from_mime(mime)
|
||||
except ValueError:
|
||||
pass
|
||||
raise ValueError(f"No MIME types matched any supported Video Codecs in '{codecs}'")
|
||||
|
||||
@staticmethod
|
||||
def from_netflix_profile(profile: str) -> Video.Codec:
|
||||
profile = profile.lower().strip()
|
||||
if profile.startswith(("h264", "playready-h264")):
|
||||
return Video.Codec.AVC
|
||||
if profile.startswith("hevc"):
|
||||
return Video.Codec.HEVC
|
||||
if profile.startswith("vp9"):
|
||||
return Video.Codec.VP9
|
||||
if profile.startswith("av1"):
|
||||
return Video.Codec.AV1
|
||||
raise ValueError(f"The Content Profile '{profile}' is not a supported Video Codec")
|
||||
|
||||
class Range(str, Enum):
|
||||
SDR = "SDR" # No Dynamic Range
|
||||
HLG = "HLG" # https://en.wikipedia.org/wiki/Hybrid_log%E2%80%93gamma
|
||||
HDR10 = "HDR10" # https://en.wikipedia.org/wiki/HDR10
|
||||
HDR10P = "HDR10+" # https://en.wikipedia.org/wiki/HDR10%2B
|
||||
DV = "DV" # https://en.wikipedia.org/wiki/Dolby_Vision
|
||||
|
||||
@staticmethod
|
||||
def from_cicp(primaries: int, transfer: int, matrix: int) -> Video.Range:
|
||||
"""
|
||||
ISO/IEC 23001-8 Coding-independent code points to Video Range.
|
||||
|
||||
Sources:
|
||||
https://www.itu.int/rec/T-REC-H.Sup19-202104-I
|
||||
"""
|
||||
|
||||
class Primaries(Enum):
|
||||
Unspecified = 0
|
||||
BT_709 = 1
|
||||
BT_601_625 = 5
|
||||
BT_601_525 = 6
|
||||
BT_2020_and_2100 = 9
|
||||
SMPTE_ST_2113_and_EG_4321 = 12 # P3D65
|
||||
|
||||
class Transfer(Enum):
|
||||
Unspecified = 0
|
||||
BT_709 = 1
|
||||
BT_601 = 6
|
||||
BT_2020 = 14
|
||||
BT_2100 = 15
|
||||
BT_2100_PQ = 16
|
||||
BT_2100_HLG = 18
|
||||
|
||||
class Matrix(Enum):
|
||||
RGB = 0
|
||||
YCbCr_BT_709 = 1
|
||||
YCbCr_BT_601_625 = 5
|
||||
YCbCr_BT_601_525 = 6
|
||||
YCbCr_BT_2020_and_2100 = 9 # YCbCr BT.2100 shares the same CP
|
||||
ICtCp_BT_2100 = 14
|
||||
|
||||
if transfer == 5:
|
||||
# While not part of any standard, it is typically used as a PAL variant of Transfer.BT_601=6.
|
||||
# i.e. where Transfer 6 would be for BT.601-NTSC and Transfer 5 would be for BT.601-PAL.
|
||||
# The codebase is currently agnostic to either, so a manual conversion to 6 is done.
|
||||
transfer = 6
|
||||
|
||||
primaries = Primaries(primaries)
|
||||
transfer = Transfer(transfer)
|
||||
matrix = Matrix(matrix)
|
||||
|
||||
# primaries and matrix does not strictly correlate to a range
|
||||
|
||||
if (primaries, transfer, matrix) == (0, 0, 0):
|
||||
return Video.Range.SDR
|
||||
elif primaries in (Primaries.BT_601_625, Primaries.BT_601_525):
|
||||
return Video.Range.SDR
|
||||
elif transfer == Transfer.BT_2100_PQ:
|
||||
return Video.Range.HDR10
|
||||
elif transfer == Transfer.BT_2100_HLG:
|
||||
return Video.Range.HLG
|
||||
else:
|
||||
return Video.Range.SDR
|
||||
|
||||
@staticmethod
|
||||
def from_m3u_range_tag(tag: str) -> Optional[Video.Range]:
|
||||
tag = (tag or "").upper().replace('"', "").strip()
|
||||
if not tag:
|
||||
return None
|
||||
if tag == "SDR":
|
||||
return Video.Range.SDR
|
||||
elif tag == "PQ":
|
||||
return Video.Range.HDR10 # technically could be any PQ-transfer range
|
||||
elif tag == "HLG":
|
||||
return Video.Range.HLG
|
||||
# for some reason there's no Dolby Vision info tag
|
||||
raise ValueError(f"The M3U Range Tag '{tag}' is not a supported Video Range")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*args: Any,
|
||||
codec: Optional[Video.Codec] = None,
|
||||
range_: Optional[Video.Range] = None,
|
||||
bitrate: Optional[Union[str, int, float]] = None,
|
||||
width: Optional[int] = None,
|
||||
height: Optional[int] = None,
|
||||
fps: Optional[Union[str, int, float]] = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""
|
||||
Create a new Video track object.
|
||||
|
||||
Parameters:
|
||||
codec: A Video.Codec enum representing the video codec.
|
||||
If not specified, MediaInfo will be used to retrieve the codec
|
||||
once the track has been downloaded.
|
||||
range_: A Video.Range enum representing the video color range.
|
||||
Defaults to SDR if not specified.
|
||||
bitrate: A number or float representing the average bandwidth in bytes/s.
|
||||
Float values are rounded up to the nearest integer.
|
||||
width: The horizontal resolution of the video.
|
||||
height: The vertical resolution of the video.
|
||||
fps: A number, float, or string representing the frames/s of the video.
|
||||
Strings may represent numbers, floats, or a fraction (num/den).
|
||||
All strings will be cast to either a number or float.
|
||||
|
||||
Note: If codec, bitrate, width, height, or fps is not specified some checks
|
||||
may be skipped or assume a value. Specifying as much information as possible
|
||||
is highly recommended.
|
||||
"""
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
if not isinstance(codec, (Video.Codec, type(None))):
|
||||
raise TypeError(f"Expected codec to be a {Video.Codec}, not {codec!r}")
|
||||
if not isinstance(range_, (Video.Range, type(None))):
|
||||
raise TypeError(f"Expected range_ to be a {Video.Range}, not {range_!r}")
|
||||
if not isinstance(bitrate, (str, int, float, type(None))):
|
||||
raise TypeError(f"Expected bitrate to be a {str}, {int}, or {float}, not {bitrate!r}")
|
||||
if not isinstance(width, (int, str, type(None))):
|
||||
raise TypeError(f"Expected width to be a {int}, not {width!r}")
|
||||
if not isinstance(height, (int, str, type(None))):
|
||||
raise TypeError(f"Expected height to be a {int}, not {height!r}")
|
||||
if not isinstance(fps, (str, int, float, type(None))):
|
||||
raise TypeError(f"Expected fps to be a {str}, {int}, or {float}, not {fps!r}")
|
||||
|
||||
self.codec = codec
|
||||
self.range = range_ or Video.Range.SDR
|
||||
|
||||
try:
|
||||
self.bitrate = int(math.ceil(float(bitrate))) if bitrate else None
|
||||
except (ValueError, TypeError) as e:
|
||||
raise ValueError(f"Expected bitrate to be a number or float, {e}")
|
||||
|
||||
try:
|
||||
self.width = int(width or 0) or None
|
||||
except ValueError as e:
|
||||
raise ValueError(f"Expected width to be a number, not {width!r}, {e}")
|
||||
|
||||
try:
|
||||
self.height = int(height or 0) or None
|
||||
except ValueError as e:
|
||||
raise ValueError(f"Expected height to be a number, not {height!r}, {e}")
|
||||
|
||||
try:
|
||||
self.fps = (FPS.parse(str(fps)) or None) if fps else None
|
||||
except Exception as e:
|
||||
raise ValueError("Expected fps to be a number, float, or a string as numerator/denominator form, " + str(e))
|
||||
|
||||
def __str__(self) -> str:
|
||||
return " | ".join(
|
||||
filter(
|
||||
bool,
|
||||
[
|
||||
"VID",
|
||||
"[" + (", ".join(filter(bool, [self.codec.value if self.codec else None, self.range.name]))) + "]",
|
||||
str(self.language),
|
||||
", ".join(
|
||||
filter(
|
||||
bool,
|
||||
[
|
||||
" @ ".join(
|
||||
filter(
|
||||
bool,
|
||||
[
|
||||
f"{self.width}x{self.height}" if self.width and self.height else None,
|
||||
f"{self.bitrate // 1000} kb/s" if self.bitrate else None,
|
||||
],
|
||||
)
|
||||
),
|
||||
f"{self.fps:.3f} FPS" if self.fps else None,
|
||||
],
|
||||
)
|
||||
),
|
||||
self.edition,
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
def change_color_range(self, range_: int) -> None:
|
||||
"""Change the Video's Color Range to Limited (0) or Full (1)."""
|
||||
if not self.path or not self.path.exists():
|
||||
raise ValueError("Cannot change the color range flag on a Video that has not been downloaded.")
|
||||
if not self.codec:
|
||||
raise ValueError("Cannot change the color range flag on a Video that has no codec specified.")
|
||||
if self.codec not in (Video.Codec.AVC, Video.Codec.HEVC):
|
||||
raise NotImplementedError(
|
||||
"Cannot change the color range flag on this Video as "
|
||||
f"it's codec, {self.codec.value}, is not yet supported."
|
||||
)
|
||||
|
||||
if not binaries.FFMPEG:
|
||||
raise EnvironmentError('FFmpeg executable "ffmpeg" was not found but is required for this call.')
|
||||
|
||||
filter_key = {Video.Codec.AVC: "h264_metadata", Video.Codec.HEVC: "hevc_metadata"}[self.codec]
|
||||
|
||||
original_path = self.path
|
||||
output_path = original_path.with_stem(f"{original_path.stem}_{['limited', 'full'][range_]}_range")
|
||||
|
||||
subprocess.run(
|
||||
[
|
||||
binaries.FFMPEG,
|
||||
"-hide_banner",
|
||||
"-loglevel",
|
||||
"panic",
|
||||
"-i",
|
||||
original_path,
|
||||
"-codec",
|
||||
"copy",
|
||||
"-bsf:v",
|
||||
f"{filter_key}=video_full_range_flag={range_}",
|
||||
str(output_path),
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
|
||||
self.path = output_path
|
||||
original_path.unlink()
|
||||
|
||||
def ccextractor(
|
||||
self, track_id: Any, out_path: Union[Path, str], language: Language, original: bool = False
|
||||
) -> Optional[Subtitle]:
|
||||
"""Return a TextTrack object representing CC track extracted by CCExtractor."""
|
||||
if not self.path:
|
||||
raise ValueError("You must download the track first.")
|
||||
|
||||
if not binaries.CCExtractor:
|
||||
raise EnvironmentError("ccextractor executable was not found.")
|
||||
|
||||
# ccextractor often fails in weird ways unless we repack
|
||||
self.repackage()
|
||||
|
||||
out_path = Path(out_path)
|
||||
|
||||
try:
|
||||
subprocess.run(
|
||||
[binaries.CCExtractor, "-trim", "-nobom", "-noru", "-ru1", "-o", out_path, self.path],
|
||||
check=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
)
|
||||
except subprocess.CalledProcessError as e:
|
||||
out_path.unlink(missing_ok=True)
|
||||
if not e.returncode == 10: # No captions found
|
||||
raise
|
||||
|
||||
if out_path.exists():
|
||||
cc_track = Subtitle(
|
||||
id_=track_id,
|
||||
url="", # doesn't need to be downloaded
|
||||
codec=Subtitle.Codec.SubRip,
|
||||
language=language,
|
||||
is_original_lang=original,
|
||||
cc=True,
|
||||
)
|
||||
cc_track.path = out_path
|
||||
return cc_track
|
||||
|
||||
return None
|
||||
|
||||
def extract_c608(self) -> list[Subtitle]:
|
||||
"""
|
||||
Extract Apple-Style c608 box (CEA-608) subtitle using ccextractor.
|
||||
|
||||
This isn't much more than a wrapper to the track.ccextractor function.
|
||||
All this does, is actually check if a c608 box exists and only if so
|
||||
does it actually call ccextractor.
|
||||
|
||||
Even though there is a possibility of more than one c608 box, only one
|
||||
can actually be extracted. Not only that but it's very possible this
|
||||
needs to be done before any decryption as the decryption may destroy
|
||||
some of the metadata.
|
||||
|
||||
TODO: Need a test file with more than one c608 box to add support for
|
||||
more than one CEA-608 extraction.
|
||||
"""
|
||||
if not self.path:
|
||||
raise ValueError("You must download the track first.")
|
||||
with self.path.open("rb") as f:
|
||||
# assuming 20KB is enough to contain the c608 box.
|
||||
# ffprobe will fail, so a c608 box check must be done.
|
||||
c608_count = len(list(get_boxes(f.read(20000), b"c608")))
|
||||
if c608_count > 0:
|
||||
# TODO: Figure out the real language, it might be different
|
||||
# CEA-608 boxes doesnt seem to carry language information :(
|
||||
# TODO: Figure out if the CC language is original lang or not.
|
||||
# Will need to figure out above first to do so.
|
||||
track_id = f"ccextractor-{self.id}"
|
||||
cc_lang = self.language
|
||||
cc_track = self.ccextractor(
|
||||
track_id=track_id,
|
||||
out_path=config.directories.temp / config.filenames.subtitle.format(id=track_id, language=cc_lang),
|
||||
language=cc_lang,
|
||||
original=False,
|
||||
)
|
||||
if not cc_track:
|
||||
return []
|
||||
return [cc_track]
|
||||
return []
|
||||
|
||||
def remove_eia_cc(self) -> bool:
|
||||
"""
|
||||
Remove EIA-CC data from Bitstream while keeping SEI data.
|
||||
|
||||
This works by removing all NAL Unit's with the Type of 6 from the bistream
|
||||
and then re-adding SEI data (effectively a new NAL Unit with just the SEI data).
|
||||
Only bitstreams with x264 encoding information is currently supported due to the
|
||||
obscurity on the MDAT mp4 box structure. Therefore, we need to use hacky regex.
|
||||
"""
|
||||
if not self.path or not self.path.exists():
|
||||
raise ValueError("Cannot clean a Track that has not been downloaded.")
|
||||
|
||||
if not binaries.FFMPEG:
|
||||
raise EnvironmentError('FFmpeg executable "ffmpeg" was not found but is required for this call.')
|
||||
|
||||
log = logging.getLogger("x264-clean")
|
||||
log.info("Removing EIA-CC from Video Track with FFMPEG")
|
||||
|
||||
with open(self.path, "rb") as f:
|
||||
file = f.read(60000)
|
||||
|
||||
x264 = re.search(rb"(.{16})(x264)", file)
|
||||
if not x264:
|
||||
log.info(" - No x264 encode settings were found, unsupported...")
|
||||
return False
|
||||
|
||||
uuid = x264.group(1).hex()
|
||||
i = file.index(b"x264")
|
||||
encoding_settings = file[i : i + file[i:].index(b"\x00")].replace(b":", rb"\\:").replace(b",", rb"\,").decode()
|
||||
|
||||
original_path = self.path
|
||||
cleaned_path = original_path.with_suffix(f".cleaned{original_path.suffix}")
|
||||
subprocess.run(
|
||||
[
|
||||
binaries.FFMPEG,
|
||||
"-hide_banner",
|
||||
"-loglevel",
|
||||
"panic",
|
||||
"-i",
|
||||
original_path,
|
||||
"-map_metadata",
|
||||
"-1",
|
||||
"-fflags",
|
||||
"bitexact",
|
||||
"-bsf:v",
|
||||
f"filter_units=remove_types=6,h264_metadata=sei_user_data={uuid}+{encoding_settings}",
|
||||
"-codec",
|
||||
"copy",
|
||||
str(cleaned_path),
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
|
||||
log.info(" + Removed")
|
||||
|
||||
self.path = cleaned_path
|
||||
original_path.unlink()
|
||||
|
||||
return True
|
||||
|
||||
|
||||
__all__ = ("Video",)
|
||||
Reference in New Issue
Block a user