From 5ec816c16359702da6dadbcc612246383817b37d Mon Sep 17 00:00:00 2001 From: Esa Kataja Date: Thu, 17 Apr 2025 12:58:07 +0300 Subject: [PATCH] Optimizations to stream classes --- src/models/streams.py | 106 ++++++++++++++++++++++++------------------ 1 file changed, 62 insertions(+), 44 deletions(-) diff --git a/src/models/streams.py b/src/models/streams.py index b45ccff..89e4582 100644 --- a/src/models/streams.py +++ b/src/models/streams.py @@ -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