Оптимизация финальной склейки

Author: Olexasha

.env

@@ -4,7 +4,8 @@ START_BATCH_TO_UPSCALE=1
 END_BATCH_TO_UPSCALE=0
 STEP_PER_BATCH=6
 FRAMES_PER_BATCH=1000
-ALLOWED_THREADS=12
+ALLOWED_CPU_THREADS=12
+ALLOWED_GPU_THREADS=2
 
 # Настройка апскейла
 MODEL_NAME=realesr-animevideov3

main.py

@@ -1,5 +1,6 @@
 import asyncio
 import os
+import re
 from datetime import datetime
 
 from src.audio.audio_handling import AudioHandler
@@ -28,23 +29,25 @@ async def main():
     start_time = datetime.now()
     print_header("запуск обработки видео")
 
-    print("\n📹 Получаем FPS исходного видео...")
     fps = await asyncio.to_thread(get_fps_accurate, ORIGINAL_VIDEO)
-    print(f"✅ FPS видео: {fps}")
-
     audio = AudioHandler()
     video = VideoHandler(fps=fps)
 
-    print("\n🔊 Извлекаем аудиодорожку...")
     await asyncio.create_task(audio.extract_audio())
     await asyncio.to_thread(extract_frames_to_batches)
-    print("✅ Кадры успешно извлечены")
 
     # Определяем диапазон батчей для обработки
     start_batch = START_BATCH_TO_UPSCALE
     end_batch = 0
     if END_BATCH_TO_UPSCALE == 0:
-        end_batch_to_upscale = len(os.listdir(INPUT_BATCHES_DIR)) - 1  # -1 для .gitkeep
+        batch_name_pattern = re.compile(r"batch_(\d+)")
+        end_batch_to_upscale = len(
+            [
+                batch
+                for batch in os.listdir(INPUT_BATCHES_DIR)
+                if batch_name_pattern.match(batch)
+            ]
+        )
     else:
         end_batch_to_upscale = END_BATCH_TO_UPSCALE
     print(f"\nВсего батчей для обработки: {end_batch_to_upscale}")
@@ -65,16 +68,16 @@ async def main():
         print(f"✅ Батчи {start_batch}-{end_batch} успешно апскейлены")
 
         batches_to_perform = [f"batch_{i}" for i in range(start_batch, end_batch + 1)]
-        short_video = await video.process_frames_to_video(batches_to_perform)
+        short_video = await video.build_short_video(batches_to_perform)
         if short_video:
+            await asyncio.to_thread(delete_frames, del_upscaled=False)
             await asyncio.to_thread(delete_frames, del_upscaled=True)
             print(f"🎥 Видео собрано: {short_video}")
             print(f"🗑️ Обработанные кадры удалены из батчей {start_batch}-{end_batch}")
 
         start_batch += STEP_PER_BATCH
 
     print_header("финальная сборка видео")
-    await asyncio.to_thread(delete_frames, del_upscaled=False)
     # Сборка финального видео из временных видео
     final_merge = video.build_final_video()
 

src/audio/audio_handling.py

@@ -9,12 +9,12 @@
     run_ffmpeg_command_with_progress,
 )
 from src.config.settings import (
-    ALLOWED_THREADS,
+    ALLOWED_CPU_THREADS,
     AUDIO_PATH,
     FINAL_VIDEO,
     ORIGINAL_VIDEO,
     RESOLUTION,
-    TMP_VIDEO,
+    TMP_VIDEO_PATH,
 )
 from src.utils.file_utils import delete_file
 from src.video.video_helpers import get_video_duration
@@ -32,7 +32,7 @@ class AudioHandler:
     def __init__(
         self,
         input_video_path: str = ORIGINAL_VIDEO,
-        merged_video_path: str = TMP_VIDEO,
+        merged_video_path: str = TMP_VIDEO_PATH,
         output_video_path: str = FINAL_VIDEO,
         audio_path: str = AUDIO_PATH,
         audio_format: str = "mp3",
@@ -52,20 +52,6 @@ def __init__(
         print(f"\tЧастота дискретизации: {self.SAMPLE_FREQ} Hz")
         print(f"\tКаналы: {self.CANALS} (стерео)")
 
-    def check_audio_extracted(self, audio_file) -> None:
-        """
-        Проверяет, было ли аудио успешно извлечено из видеофайла.
-        Возвращает True, если аудио существует, иначе False.
-        """
-        if audio_file and os.path.exists(audio_file):
-            print(f"✅ Аудио успешно извлечено: {audio_file}\n")
-            self.audio_path = audio_file
-        else:
-            print("⚠️ Аудио не найдено или не было извлечено.")
-            raise FileNotFoundError(
-                "Аудиофайл не найден. Проверьте, было ли аудио успешно извлечено."
-            )
-
     async def extract_audio(self) -> Optional[str]:
         """
         Извлекает аудио из видеофайла и сохраняет его как отдельный аудиофайл.
@@ -86,7 +72,7 @@ def __sync_extract():
                 "-ac", self.CANALS,
                 "-b:a", self.BITRATE,
                 "-progress", "-",
-                "-threads", str(ALLOWED_THREADS),
+                "-threads", str(ALLOWED_CPU_THREADS),
                 "-loglevel", "error",
                 audio_file,
             ]
@@ -97,7 +83,8 @@ def __sync_extract():
             except subprocess.CalledProcessError as e:
                 print(f"🚨 Ошибка при извлечении аудио: {e}")
                 return None
-            self.check_audio_extracted(audio_file)
+            self.__check_audio_extracted(audio_file)
+            return None
 
         loop = asyncio.get_event_loop()
         return await loop.run_in_executor(None, __sync_extract)
@@ -119,7 +106,7 @@ def __sync_insert():
             print(f"\tКодек аудио: исходный ({self.audio_format})")
             print(f"\tДлительность видео: {duration:.2f} сек")
             print(f"\tFPS видео: {fps}")
-            print(f"\tПотоков: {ALLOWED_THREADS}")
+            print(f"\tПотоков: {ALLOWED_CPU_THREADS}")
             print(f"\tРазрешение: {self.resolution}")
 
             cmd = [
@@ -130,7 +117,7 @@ def __sync_insert():
                 "-map", "0:v:0",
                 "-map", "1:a:0",
                 "-shortest", "-progress", "-",
-                "-threads", str(ALLOWED_THREADS),
+                "-threads", str(ALLOWED_CPU_THREADS),
                 "-nostats", "-loglevel", "error",
                 self.out_video_path,
             ]
@@ -147,6 +134,20 @@ def __sync_insert():
         loop = asyncio.get_event_loop()
         await loop.run_in_executor(None, __sync_insert)
 
+    def __check_audio_extracted(self, audio_file) -> None:
+        """
+        Проверяет, было ли аудио успешно извлечено из видеофайла.
+        Возвращает True, если аудио существует, иначе False.
+        """
+        if audio_file and os.path.exists(audio_file):
+            print(f"✅ Аудио успешно извлечено: {audio_file}\n")
+            self.audio_path = audio_file
+        else:
+            print("⚠️ Аудио не найдено или не было извлечено.")
+            raise FileNotFoundError(
+                "Аудиофайл не найден. Проверьте, было ли аудио успешно извлечено."
+            )
+
     def __str__(self):
         return (
             f"AudioHandler(in_video_path={self.in_video_path}, "

src/config/settings.py

@@ -8,7 +8,7 @@
 # Пути к файлам
 ORIGINAL_VIDEO = os.path.join(ROOT_DIR, "data", "input_video", "naruto_test.mp4")
 AUDIO_PATH = os.path.join(ROOT_DIR, "data", "audio")
-TMP_VIDEO = os.path.join(ROOT_DIR, "data", "tmp_video", "final_merge.mp4")
+TMP_VIDEO_PATH = os.path.join(ROOT_DIR, "data", "tmp_video")
 FINAL_VIDEO = os.path.join(ROOT_DIR, "data", "output_video", "final_video.mp4")
 BATCH_VIDEO_PATH = os.path.join(ROOT_DIR, "data", "video_batches")
 INPUT_BATCHES_DIR = os.path.join(ROOT_DIR, "data", "default_frame_batches")
@@ -20,7 +20,8 @@
 END_BATCH_TO_UPSCALE = int(os.getenv("END_BATCH_TO_UPSCALE", 0))
 STEP_PER_BATCH = int(os.getenv("STEP_PER_BATCH", 6))
 FRAMES_PER_BATCH = int(os.getenv("FRAMES_PER_BATCH", 1000))
-ALLOWED_THREADS = int(os.getenv("ALLOWED_THREADS", 6))
+ALLOWED_CPU_THREADS = int(os.getenv("ALLOWED_THREADS", 6))
+ALLowed_GPU_THREADS = int(os.getenv("ALLOWED_GPU_THREADS", 1))
 
 # Настройка апскейла
 MODEL_DIR = os.path.join(ROOT_DIR, "src", "utils", "realesrgan", "models")

src/frames/frames_helpers.py

@@ -8,7 +8,7 @@
 from tqdm import tqdm
 
 from src.config.settings import (
-    ALLOWED_THREADS,
+    ALLOWED_CPU_THREADS,
     FRAMES_PER_BATCH,
     INPUT_BATCHES_DIR,
     ORIGINAL_VIDEO,
@@ -80,7 +80,7 @@ def extract_frames_to_batches(
         colour="green",
         file=sys.stdout,
     ) as pbar:
-        with ThreadPoolExecutor(max_workers=ALLOWED_THREADS) as executor:
+        with ThreadPoolExecutor(max_workers=ALLOWED_CPU_THREADS) as executor:
             current_batch_dir = make_default_batch_dir(output_dir)
 
             while True:

src/frames/upscale.py

@@ -10,7 +10,7 @@
 from tqdm import tqdm
 
 from src.config.settings import (
-    ALLOWED_THREADS,
+    ALLOWED_CPU_THREADS,
     INPUT_BATCHES_DIR,
     MODEL_DIR,
     MODEL_NAME,
@@ -118,7 +118,7 @@ async def upscale_batches(start_batch: int, end_batch: int):
         # Запускаем обработку батчей с использованием ProcessPoolExecutor
         loop = asyncio.get_event_loop()
 
-        with ProcessPoolExecutor(max_workers=ALLOWED_THREADS) as executor:
+        with ProcessPoolExecutor(max_workers=ALLOWED_CPU_THREADS) as executor:
             tasks = [
                 loop.run_in_executor(executor, _upscale, batch_num)
                 for batch_num in batches_range

src/video/video_exceptions.py

@@ -0,0 +1,17 @@
+class VideoDoesNotExist(Exception):
+    def __init__(self, video_path: str):
+        super().__init__(
+            f"Видео '{video_path}' не найдено. Проверьте правильность пути."
+        )
+
+
+class VideoMergingError(Exception):
+    def __init__(self, error):
+        super().__init__(f"Ошибка при объединении видео.\nПричина: {error}")
+
+
+class VideoReadFrameError(Exception):
+    def __init__(self, video_path: str):
+        super().__init__(
+            f"Не удалось прочитать кадр {1} из видео '{video_path}'."
+        )