Add Base functionality

2024-11-14 23:46:13 +02:00
parent c298763049
commit cf3b7bae5d
8 changed files with 383 additions and 1 deletions
@@ -0,0 +1,134 @@
+from pydantic import BaseModel, Field
+from typing import Optional
+from datetime import datetime
+from pathlib import Path
+
+from lib.fs import sanitize_filename
+
+
+class Tags(BaseModel):
+    title: Optional[str] = Field(None)
+    language: Optional[str] = Field(None)
+
+
+class Disposition(BaseModel):
+    default: bool = Field(False)
+    comment: bool = Field(False)
+
+
+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)
+    disposition: Disposition
+    tags: Tags
+
+    @property
+    def fps(self) -> float:
+        frame_numbers = [int(x) for x in self.avg_frame_rate.split("/")]
+        return frame_numbers[0] / frame_numbers[1]
+
+    def cmd_flags(self) -> list[str]:
+        return [
+            "-map",
+            f"0:{self.index}",
+            "-c:v",
+            "libsvtav1",
+            "-crf",
+            "30",
+            "-preset",
+            "6",
+            "-svtav1-params",
+            "tune=0:film-grain=20",
+            "-vf",
+            "'scale=iw*sar:ih,setsar=1,scale=-2:ih:lanczos'",
+            "-g",
+            f"{round(self.fps)*5}",
+        ]
+
+
+class AudioStream(BaseModel):
+    index: int
+    codec_name: str
+    sample_fmt: str
+    channel_layout: str
+    disposition: Disposition
+    tags: Tags
+
+    def cmd_flags(self, index: int) -> list[str]:
+        side_loaded = []
+        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'"]
+        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}'",
+            ]
+            + default_disposition
+        )
+
+
+class SubtitleStream(BaseModel):
+    index: int
+    codec_name: str
+    disposition: Disposition
+    tags: Tags
+
+    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}",
+        ]
+
+
+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([])
+    allow_subtitle: bool = Field(True)
+
+    def cmd(self) -> list[str]:
+        cmd = [
+            "ffmpeg",
+            "-hide_banner",
+            "-loglevel",
+            "error",
+            "-i",
+            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}",
+        ]
+        for stream in self.video_streams:
+            cmd.extend(stream.cmd_flags())
+        for i, stream in enumerate(self.audio_streams):
+            cmd.extend(stream.cmd_flags(i))
+        if self.allow_subtitle:
+            for i, stream in enumerate(self.subtitle_streams):
+                cmd.extend(stream.cmd_flags(i))
+
+        output_file = (
+            self.video_path.absolute().parent
+            / f"{sanitize_filename(self.title)}.{self.date_released.year}.mkv"
+        )
+
+        cmd.extend([str(output_file)])
+
+        return cmd