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 typing import Optional
from enum import Enum
from datetime import datetime
from pathlib import Path
from functools import cached_property
from lib.fs import sanitize_filename
class OriginalMediaType(str, Enum):
DVD = "DVD"
BluRay = "BluRay"
TVRip = "TVRip"
WebRip = "WebRip"
Other = "Other"
class Tags(BaseModel):
title: Optional[str] = Field(None)
language: Optional[str] = Field(None)
title: str | None = Field(None)
language: str | None = Field(None)
class Disposition(BaseModel):
@@ -18,24 +25,31 @@ class Disposition(BaseModel):
class VideoStream(BaseModel):
index: int = Field(None)
width: int = Field(None)
height: int = Field(None)
avg_frame_rate: str = Field(None, alias="avg_frame_rate")
codec_name: str = Field(None)
index: int | None = Field(None)
width: int | None = Field(None)
height: int | None = Field(None)
avg_frame_rate: str | None = Field(None, alias="avg_frame_rate")
codec_name: str | None = Field(None)
disposition: Disposition
tags: Tags | None = Field(None)
preset: int = Field(2)
crf: int = Field(30)
original_media_type: str = Field(None)
original_media_type: str | None = Field(None)
film_grain: int = Field(20)
@cached_property
def fps(self) -> float:
frame_numbers = [int(x) for x in self.avg_frame_rate.split("/")]
return frame_numbers[0] / frame_numbers[1]
"""Calculate frames per second from avg_frame_rate string."""
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]:
"""Generate ffmpeg command flags for this video stream."""
return [
"-map",
f"0:{self.index}",
@@ -50,7 +64,7 @@ class VideoStream(BaseModel):
"-vf",
"'scale=iw*sar:ih,setsar=1,scale=-2:ih:lanczos'",
"-g",
f"{round(self.fps) * 5}",
str(round(self.fps) * 5),
]
@@ -63,20 +77,24 @@ class AudioStream(BaseModel):
tags: Tags | None = Field(None)
def cmd_flags(self, index: int) -> list[str]:
"""Generate ffmpeg command flags for this audio stream."""
side_loaded = []
if "side" in self.channel_layout:
side_loaded = [f"-filter:a:{index}", "'channelmap=channel_layout=5.1'"]
default_disposition = (
[f"-disposition:a:{index}", "default"] if self.disposition.default else []
)
if "side" in self.channel_layout:
side_loaded = [f"-filter:a:{index}", "'channelmap=channel_layout=5.1'"]
language = getattr(self.tags, 'language', None) if self.tags else None
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 (
["-map", f"0:{self.index}", f"-c:a:{index}", "libopus"]
+ side_loaded
+ [f"-metadata:s:a:{index}", f"language={self.tags.language}"]
+ [
f"-metadata:s:a:{index}",
f"title='{self.tags.title}'",
]
+ meta_flags
+ default_disposition
)
@@ -88,23 +106,27 @@ class SubtitleStream(BaseModel):
tags: Tags | None = Field(None)
def cmd_flags(self, index: int) -> list[str]:
return ["-map", f"0:s:{index}", f"-c:s:{index}", "copy"] + [
f"-metadata:s:s:{index}",
f"language={self.tags.language}",
]
"""Generate ffmpeg command flags for this subtitle stream."""
language = getattr(self.tags, 'language', None) if self.tags else None
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):
video_path: Path
title: str = Field(None)
date_released: Optional[datetime] = Field(None)
imdb: Optional[str] = Field(None)
video_streams: list[VideoStream] = Field([])
audio_streams: list[AudioStream] = Field([])
subtitle_streams: list[SubtitleStream] = Field([])
title: str | None = Field(None)
date_released: datetime | None = Field(None)
imdb: str | None = Field(None)
video_streams: list[VideoStream] = Field(default_factory=list)
audio_streams: list[AudioStream] = Field(default_factory=list)
subtitle_streams: list[SubtitleStream] = Field(default_factory=list)
allow_subtitle: bool = Field(True)
original_media_type: OriginalMediaType = Field(OriginalMediaType.DVD)
def cmd(self) -> list[str]:
"""Generate the full ffmpeg command for this video file."""
cmd = [
"ffmpeg",
"-hide_banner",
@@ -114,17 +136,15 @@ class VideoFile(BaseModel):
f"'{str(self.video_path.absolute())}'",
"-map_metadata",
"-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:
cmd.extend(stream.cmd_flags())
for i, stream in enumerate(self.audio_streams):
@@ -132,12 +152,10 @@ class VideoFile(BaseModel):
if self.allow_subtitle:
for i, stream in enumerate(self.subtitle_streams):
cmd.extend(stream.cmd_flags(i))
year = self.date_released.year if self.date_released else "unknown"
output_file = (
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.extend([f"'{str(output_file)}'"])
cmd.append(f"'{str(output_file)}'")
return cmd