From 61f8e0abe12d32aa5be20a67bc95d50b4568c557 Mon Sep 17 00:00:00 2001 From: Viner Abubakirov Date: Fri, 3 Apr 2026 17:03:10 +0500 Subject: [PATCH] =?UTF-8?q?=D0=9F=D0=BE=D0=BC=D0=B5=D0=BD=D1=8F=D0=BB,=20?= =?UTF-8?q?=D0=BC=D0=B5=D1=82=D0=BE=D0=B4=20=D1=84=D0=BE=D1=80=D0=BC=D0=B8?= =?UTF-8?q?=D1=80=D0=BE=D0=B2=D0=B0=D0=BD=D0=B8=D1=8F=20=D0=B2=D0=B8=D0=B4?= =?UTF-8?q?=D0=B5=D0=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 42 +++++++++++++++++++----------------- src/utils/video.py | 54 ++++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 74 insertions(+), 22 deletions(-) diff --git a/main.py b/main.py index 35ba949..170171b 100644 --- a/main.py +++ b/main.py @@ -98,9 +98,8 @@ class InterpolationPipeline: def run(self, video_path: Path, output_video: str): prev_frames = tuple() - interpolated_frames = [] + interpolated_frames: list["np.ndarray"] = [] part = 0 - source_frame_length = 0 chunk_seconds = 10 length = self.video_maker.get_video_duration(video_path) last_part_seconds = 1 if length % chunk_seconds else 0 @@ -108,6 +107,7 @@ class InterpolationPipeline: fps = self.video_maker.get_fps(video_path) logging.info(f"Video FPS: {fps}") fps *= 2 # Doubling FPS + width, height = self.video_maker.get_size(video_path) for frames in self.video_maker.video_to_frames_generator( video_path, self.fs.frames_path, chunk_seconds ): @@ -117,13 +117,10 @@ class InterpolationPipeline: img2 = frames[0] img1_2 = self.interpolator.interpolate(img1, img2) interpolated_frames.append(img1_2) - self.fs.clear_directory(self.fs.moved_path) - self._save_images(prev_frames, interpolated_frames) - - self._merge_frames_to_video( - self.fs.video_part_path / f"video_{part:08d}.mp4", - fps, - source_frame_length, + generator = self._frame_generator(prev_frames, interpolated_frames) + part_path = self.fs.video_part_path / f"video_{part:08d}.mp4" + self.video_maker.images_to_video_pipeline( + generator, part_path, width, height, fps ) interpolated_frames = [] logging.info(f"Finished processing part {part:08d}") @@ -136,17 +133,13 @@ class InterpolationPipeline: img2 = frames[i + 1] img1_2 = self.interpolator.interpolate(img1, img2) interpolated_frames.append(img1_2) - source_frame_length = len(frames) prev_frames = frames - self.fs.clear_directory(self.fs.moved_path) - self._save_images(prev_frames, interpolated_frames) - self._merge_frames_to_video( - self.fs.video_part_path / f"video_{part:08d}.mp4", - fps, - source_frame_length, + generator = self._frame_generator(prev_frames, interpolated_frames) + part_path = self.fs.video_part_path / f"video_{part:08d}.mp4" + self.video_maker.images_to_video_pipeline( + generator, part_path, width, height, fps ) - self.fs.clear_directory(self.fs.moved_path) logging.info(f"Finished processing part {part:08d}") self._merge_video_parts(self.fs.output_path / output_video) logging.info( @@ -159,6 +152,7 @@ class InterpolationPipeline: interpolated: list["np.ndarray"], ): logging.info("Saving images...") + self.fs.clear_directory(self.fs.moved_path) index = 0 for i, frame in enumerate(source): name = self.fs.moved_path / f"img_{index:08d}.png" @@ -170,15 +164,23 @@ class InterpolationPipeline: imwrite(name, interpolated[i]) logging.info("Success...") - def _merge_frames_to_video( - self, output_video: Path, fps: float, source_frame_length: int = 0 - ): + def _merge_frames_to_video(self, output_video: Path, fps: float): self.video_maker.images_to_video(self.fs.moved_path, output_video, fps) def _merge_video_parts(self, output_video: Path): self.video_maker.concatenate_videos(self.fs.video_part_path, output_video) self.fs.clear_directory(self.fs.video_part_path) + def _frame_generator( + self, + source: tuple["np.ndarray", ...], + interpolated: list["np.ndarray"], + ): + for i, frame in enumerate(source): + yield frame + if i < len(interpolated): + yield interpolated[i] + def runner( base_path: Path, diff --git a/src/utils/video.py b/src/utils/video.py index 5084fda..9ca64bd 100644 --- a/src/utils/video.py +++ b/src/utils/video.py @@ -2,7 +2,7 @@ import os import logging import subprocess from pathlib import Path -from typing import Generator +from typing import Generator, Iterable import cv2 import numpy as np @@ -67,7 +67,13 @@ class VideoMaker: def run_command(self, cmd: str) -> int: try: - subprocess.run(cmd, shell=True, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + subprocess.run( + cmd, + shell=True, + check=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) return 0 except subprocess.CalledProcessError as e: logging.error(f"Command failed with error: {e}") @@ -95,3 +101,47 @@ class VideoMaker: return paths.append(frame) yield tuple(paths) + + def images_to_video_pipeline( + self, + frames: Iterable[np.ndarray], + output_path: Path, + width: int, + height: int, + fps: float, + ): + pipeline = subprocess.Popen( + [ + "ffmpeg", + "-y", + "-f", "rawvideo", + "-vcodec", "rawvideo", + "-pix_fmt", "bgr24", + "-s", f"{width}x{height}", + "-r", str(fps), + "-i", "-", + "-an", + "-vcodec", "libx264", + "-pix_fmt", "yuv420p", + str(output_path), + ], + stdin=subprocess.PIPE, + stderr=subprocess.DEVNULL + ) + if pipeline.stdin is None: + raise Exception("STDIN closed") + for frame in frames: + pipeline.stdin.write(frame.tobytes()) + + pipeline.stdin.close() + pipeline.wait() + + def get_size(self, video_path: Path) -> tuple[int, int]: + cap = cv2.VideoCapture(str(video_path)) + if not cap.isOpened(): + raise ValueError(f"Cannot open video: {video_path}") + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + cap.release() + return width, height