Optimizations to stream classes

This commit is contained in:
Esa Kataja
2025-04-17 12:58:07 +03:00
parent c7a381c4b9
commit 5ec816c163
+62 -44
View File
@@ -1,15 +1,22 @@
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from typing import Optional from enum import Enum
from datetime import datetime from datetime import datetime
from pathlib import Path from pathlib import Path
from functools import cached_property from functools import cached_property
from lib.fs import sanitize_filename from lib.fs import sanitize_filename
class OriginalMediaType(str, Enum):
DVD = "DVD"
BluRay = "BluRay"
TVRip = "TVRip"
WebRip = "WebRip"
Other = "Other"
class Tags(BaseModel): class Tags(BaseModel):
title: Optional[str] = Field(None) title: str | None = Field(None)
language: Optional[str] = Field(None) language: str | None = Field(None)
class Disposition(BaseModel): class Disposition(BaseModel):
@@ -18,24 +25,31 @@ class Disposition(BaseModel):
class VideoStream(BaseModel): class VideoStream(BaseModel):
index: int = Field(None) index: int | None = Field(None)
width: int = Field(None) width: int | None = Field(None)
height: int = Field(None) height: int | None = Field(None)
avg_frame_rate: str = Field(None, alias="avg_frame_rate") avg_frame_rate: str | None = Field(None, alias="avg_frame_rate")
codec_name: str = Field(None) codec_name: str | None = Field(None)
disposition: Disposition disposition: Disposition
tags: Tags | None = Field(None) tags: Tags | None = Field(None)
preset: int = Field(2) preset: int = Field(2)
crf: int = Field(30) crf: int = Field(30)
original_media_type: str = Field(None) original_media_type: str | None = Field(None)
film_grain: int = Field(20) film_grain: int = Field(20)
@cached_property @cached_property
def fps(self) -> float: def fps(self) -> float:
frame_numbers = [int(x) for x in self.avg_frame_rate.split("/")] """Calculate frames per second from avg_frame_rate string."""
return frame_numbers[0] / frame_numbers[1] try:
if not self.avg_frame_rate:
return 0.0
num, denom = (int(x) for x in self.avg_frame_rate.split("/"))
return num / denom if denom else 0.0
except Exception:
return 0.0
def cmd_flags(self) -> list[str]: def cmd_flags(self) -> list[str]:
"""Generate ffmpeg command flags for this video stream."""
return [ return [
"-map", "-map",
f"0:{self.index}", f"0:{self.index}",
@@ -50,7 +64,7 @@ class VideoStream(BaseModel):
"-vf", "-vf",
"'scale=iw*sar:ih,setsar=1,scale=-2:ih:lanczos'", "'scale=iw*sar:ih,setsar=1,scale=-2:ih:lanczos'",
"-g", "-g",
f"{round(self.fps) * 5}", str(round(self.fps) * 5),
] ]
@@ -63,20 +77,24 @@ class AudioStream(BaseModel):
tags: Tags | None = Field(None) tags: Tags | None = Field(None)
def cmd_flags(self, index: int) -> list[str]: def cmd_flags(self, index: int) -> list[str]:
"""Generate ffmpeg command flags for this audio stream."""
side_loaded = [] side_loaded = []
if "side" in self.channel_layout:
side_loaded = [f"-filter:a:{index}", "'channelmap=channel_layout=5.1'"]
default_disposition = ( default_disposition = (
[f"-disposition:a:{index}", "default"] if self.disposition.default else [] [f"-disposition:a:{index}", "default"] if self.disposition.default else []
) )
if "side" in self.channel_layout: language = getattr(self.tags, 'language', None) if self.tags else None
side_loaded = [f"-filter:a:{index}", "'channelmap=channel_layout=5.1'"] title = getattr(self.tags, 'title', None) if self.tags else None
meta_flags = []
if language:
meta_flags += [f"-metadata:s:a:{index}", f"language={language}"]
if title:
meta_flags += [f"-metadata:s:a:{index}", f"title='{title}'"]
return ( return (
["-map", f"0:{self.index}", f"-c:a:{index}", "libopus"] ["-map", f"0:{self.index}", f"-c:a:{index}", "libopus"]
+ side_loaded + side_loaded
+ [f"-metadata:s:a:{index}", f"language={self.tags.language}"] + meta_flags
+ [
f"-metadata:s:a:{index}",
f"title='{self.tags.title}'",
]
+ default_disposition + default_disposition
) )
@@ -88,23 +106,27 @@ class SubtitleStream(BaseModel):
tags: Tags | None = Field(None) tags: Tags | None = Field(None)
def cmd_flags(self, index: int) -> list[str]: def cmd_flags(self, index: int) -> list[str]:
return ["-map", f"0:s:{index}", f"-c:s:{index}", "copy"] + [ """Generate ffmpeg command flags for this subtitle stream."""
f"-metadata:s:s:{index}", language = getattr(self.tags, 'language', None) if self.tags else None
f"language={self.tags.language}", meta_flags = []
] if language:
meta_flags += [f"-metadata:s:s:{index}", f"language={language}"]
return ["-map", f"0:s:{index}", f"-c:s:{index}", "copy"] + meta_flags
class VideoFile(BaseModel): class VideoFile(BaseModel):
video_path: Path video_path: Path
title: str = Field(None) title: str | None = Field(None)
date_released: Optional[datetime] = Field(None) date_released: datetime | None = Field(None)
imdb: Optional[str] = Field(None) imdb: str | None = Field(None)
video_streams: list[VideoStream] = Field([]) video_streams: list[VideoStream] = Field(default_factory=list)
audio_streams: list[AudioStream] = Field([]) audio_streams: list[AudioStream] = Field(default_factory=list)
subtitle_streams: list[SubtitleStream] = Field([]) subtitle_streams: list[SubtitleStream] = Field(default_factory=list)
allow_subtitle: bool = Field(True) allow_subtitle: bool = Field(True)
original_media_type: OriginalMediaType = Field(OriginalMediaType.DVD)
def cmd(self) -> list[str]: def cmd(self) -> list[str]:
"""Generate the full ffmpeg command for this video file."""
cmd = [ cmd = [
"ffmpeg", "ffmpeg",
"-hide_banner", "-hide_banner",
@@ -114,17 +136,15 @@ class VideoFile(BaseModel):
f"'{str(self.video_path.absolute())}'", f"'{str(self.video_path.absolute())}'",
"-map_metadata", "-map_metadata",
"-1", "-1",
"-metadata",
f"title='{self.title}'",
"-metadata",
f"date_released={self.date_released.strftime('%Y-%m-%d')}",
"-metadata",
f"imdb={self.imdb}",
"-metadata",
"original_media_type=DVD",
"-pix_fmt",
"yuv420p10le",
] ]
if self.title:
cmd += ["-metadata", f"title='{self.title}'"]
if self.date_released:
cmd += ["-metadata", f"date_released={self.date_released.strftime('%Y-%m-%d')}"]
if self.imdb:
cmd += ["-metadata", f"imdb={self.imdb}"]
cmd += ["-metadata", f"original_media_type={self.original_media_type.value}"]
cmd += ["-pix_fmt", "yuv420p10le"]
for stream in self.video_streams: for stream in self.video_streams:
cmd.extend(stream.cmd_flags()) cmd.extend(stream.cmd_flags())
for i, stream in enumerate(self.audio_streams): for i, stream in enumerate(self.audio_streams):
@@ -132,12 +152,10 @@ class VideoFile(BaseModel):
if self.allow_subtitle: if self.allow_subtitle:
for i, stream in enumerate(self.subtitle_streams): for i, stream in enumerate(self.subtitle_streams):
cmd.extend(stream.cmd_flags(i)) cmd.extend(stream.cmd_flags(i))
year = self.date_released.year if self.date_released else "unknown"
output_file = ( output_file = (
self.video_path.absolute().parent self.video_path.absolute().parent
/ f"{sanitize_filename(self.title)}.{self.date_released.year}.mkv" / f"{sanitize_filename(self.title or 'untitled')}.{year}.mkv"
) )
cmd.append(f"'{str(output_file)}'")
cmd.extend([f"'{str(output_file)}'"])
return cmd return cmd