| from __future__ import annotations |
|
|
| import os |
| from dataclasses import dataclass |
|
|
| from . import media_io as media |
| from . import video_buffers as video |
|
|
|
|
| @dataclass |
| class MuxSession: |
| mux_process: object | None = None |
| output_path_for_write: str = "" |
| video_only_output_path: str = "" |
| reserved_metadata_path: str | None = None |
| stopped: bool = False |
| mux_finished: bool = False |
| owns_output_path_for_write: bool = False |
| owns_video_only_output: bool = False |
|
|
| def set_output_path(self, output_path: str) -> None: |
| self.output_path_for_write = output_path |
| self.video_only_output_path = "" |
|
|
| def ensure_started( |
| self, |
| *, |
| server_config: dict, |
| ffmpeg_path: str, |
| process_is_hdr: bool, |
| use_live_av_mux: bool, |
| output_container: str, |
| source_path: str, |
| exact_start_seconds: float, |
| selected_audio_track: int | None, |
| resolved_width: int, |
| resolved_height: int, |
| fps_float: float, |
| source_audio_duration_seconds: float | None = None, |
| ) -> None: |
| if self.mux_process is not None: |
| return |
| self.reserved_metadata_path = media.create_reserved_metadata_file(self.output_path_for_write) |
| self.owns_output_path_for_write = use_live_av_mux |
| self.owns_video_only_output = not use_live_av_mux |
| if self.owns_video_only_output: |
| self.video_only_output_path = media.reserve_video_only_output_path(self.output_path_for_write) |
| if process_is_hdr: |
| hdr_video_crf = server_config.get("hdr_video_crf", 8) |
| self.mux_process = media.start_hdr_av_mux_process(ffmpeg_path, self.output_path_for_write, resolved_width, resolved_height, fps_float, hdr_video_crf, output_container, source_path, exact_start_seconds, selected_audio_track, self.reserved_metadata_path, source_audio_duration_seconds) if use_live_av_mux else media.start_hdr_video_mux_process(ffmpeg_path, self.video_only_output_path, resolved_width, resolved_height, fps_float, hdr_video_crf, output_container, self.reserved_metadata_path) |
| return |
| video_codec = server_config.get("video_output_codec", "libx264_8") |
| self.mux_process = media.start_av_mux_process(ffmpeg_path, self.output_path_for_write, resolved_width, resolved_height, fps_float, video_codec, output_container, source_path, exact_start_seconds, selected_audio_track, self.reserved_metadata_path, source_audio_duration_seconds) if use_live_av_mux else media.start_video_mux_process(ffmpeg_path, self.video_only_output_path, resolved_width, resolved_height, fps_float, video_codec, output_container, self.reserved_metadata_path) |
|
|
| def write_chunk(self, *, process_is_hdr: bool, video_tensor_hdr, video_tensor_uint8, start_frame: int, frame_count: int): |
| if process_is_hdr and video_tensor_hdr is not None: |
| return video.write_hdr_video_chunk(self.mux_process, video_tensor_hdr, start_frame=start_frame, frame_count=frame_count) |
| return video.write_video_chunk(self.mux_process, video_tensor_uint8, start_frame=start_frame, frame_count=frame_count) |
|
|
| def finalize(self) -> tuple[int, str, bool]: |
| return_code, stderr, forced_termination = media.finalize_mux_process(self.mux_process) |
| self.mux_finished = True |
| return return_code, stderr, forced_termination |
|
|
| def cleanup_partial_outputs(self) -> None: |
| if self.mux_process is not None and not self.mux_finished and self.mux_process.poll() is None: |
| try: |
| media.finalize_mux_process(self.mux_process) |
| except Exception: |
| pass |
| if self.owns_output_path_for_write and self.mux_process is not None and not self.stopped and self.mux_process.returncode not in (0, None) and os.path.isfile(self.output_path_for_write): |
| media.delete_file_if_exists(self.output_path_for_write, label="continuation output") |
| if self.owns_video_only_output: |
| media.delete_file_if_exists(self.video_only_output_path, label="video-only output") |
| media.delete_file_if_exists(self.reserved_metadata_path, label="reserved metadata file") |
|
|
