Skip to content

Commit 2f45f41

Browse files
OlexashaOlexasha
committed
Асинхронный апскейл фреймов и общий логич. прогресс
1 parent 37276de commit 2f45f41

File tree

4 files changed

+213
-253
lines changed

4 files changed

+213
-253
lines changed

‎main.py

+107-164
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import glob
22
import os
33
import shutil
4-
import subprocess
54
import asyncio
65
import ffmpeg
76
from moviepy.editor import VideoFileClip, concatenate_videoclips
@@ -10,9 +9,7 @@
109
BATCH_VIDEO_PATH,
1110
END_BATCH_TO_UPSCALE,
1211
FINAL_VIDEO,
13-
FRAMES_PER_BATCH,
1412
INPUT_BATCHES_DIR,
15-
MERGE_ALL_BATCHES_TO_VIDEO,
1613
ORIGINAL_VIDEO,
1714
OUTPUT_BATCHES_DIR,
1815
START_BATCH_TO_UPSCALE,
@@ -23,167 +20,113 @@
2320
from src.video_processing.frames import get_fps_accurate, extract_frames_to_batches
2421
from src.video_processing.upscale import upscale_batches
2522

26-
# def upscale_batches(start_batch, end_batch):
27-
# """Запускает скрипт для улучшения батчей в заданном диапазоне."""
28-
# current_dir = os.getcwd()
29-
# cmd = [
30-
# "/bin/bash",
31-
# f"{current_dir}/enhance_frames.sh",
32-
# str(start_batch),
33-
# str(end_batch),
34-
# ]
35-
# process = subprocess.Popen(
36-
# cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
37-
# )
38-
# for line in process.stdout:
39-
# print(line, end="")
40-
# process.wait()
41-
# if process.returncode == 0:
42-
# print(f"Батчи с {start_batch} по {end_batch} успешно обработаны.")
43-
# else:
44-
# print(
45-
# f"Ошибка при обработке батчей {start_batch}-{end_batch}: {process.stderr.read()}"
46-
# )
47-
#
48-
#
49-
# def merge_frames_to_video(batches_list, counter_of_video_batches):
50-
# """Собирает обработанные фреймы в видео."""
51-
# try:
52-
# batches_and_frames_list = list()
53-
# for batch in batches_list:
54-
# batch_path = f"{OUTPUT_BATCHES_DIR}/{batch}"
55-
# batches_and_frames_list.append(
56-
# sorted(glob.glob(os.path.join(batch_path, "frame*.jpg")))
57-
# )
58-
# frames_list = sum(batches_and_frames_list, [])
59-
#
60-
# # файл чтобы скормить пути фреймов ffmpeg
61-
# frames_file_path = os.path.join("./", "frames_list.txt")
62-
# with open(frames_file_path, "w") as f:
63-
# for frame in frames_list:
64-
# f.write(f"file '{frame}'\n")
65-
#
66-
# batch_video = os.path.join(
67-
# BATCH_VIDEO_PATH, f"batch_{counter_of_video_batches}.mp4"
68-
# )
69-
# ffmpeg.input(frames_file_path, format="concat", safe=0).output(
70-
# batch_video, vcodec="libx264", pix_fmt="yuv420p", crf=18, r=FPS
71-
# ).run()
72-
# os.remove(frames_file_path)
73-
# print(f"Партия {batches_list} собрана в видео.")
74-
# return batch_video
75-
# except Exception as e:
76-
# print(f"Ошибка сборки видео из фреймов: {e}")
77-
# return None
78-
#
79-
#
80-
# def delete_upscaled_frames(del_only_dirs=True):
81-
# """Удаляет обработанные фреймы."""
82-
# file_paths = glob.glob(os.path.join(OUTPUT_BATCHES_DIR, "*"))
83-
# for file_path in file_paths:
84-
# try:
85-
# if os.path.isdir(file_path):
86-
# shutil.rmtree(file_path)
87-
# if (
88-
# not del_only_dirs
89-
# and os.path.isfile(file_path)
90-
# or os.path.islink(file_path)
91-
# ):
92-
# os.remove(file_path)
93-
# except Exception as e:
94-
# print(f"Не удалось удалить {file_path}. Причина: {e}")
95-
#
96-
#
97-
# def combine_videos(tmp_final_video, batch_video):
98-
# """Объединяет партии видео в одно целое."""
99-
# try:
100-
# if os.path.isfile(tmp_final_video):
101-
# first_clip = VideoFileClip(tmp_final_video)
102-
# second_clip = VideoFileClip(batch_video)
103-
# combined_clip = concatenate_videoclips([first_clip, second_clip])
104-
# combined_clip.write_videofile(FINAL_VIDEO, codec="libx264", fps=FPS)
105-
# print(f"Объединение видеофайлов завершено.")
106-
#
107-
# first_clip.close()
108-
# second_clip.close()
109-
# combined_clip.close()
110-
# os.unlink(batch_video)
111-
# os.unlink(tmp_final_video)
112-
# os.rename(FINAL_VIDEO, tmp_final_video)
113-
# else:
114-
# shutil.move(batch_video, tmp_final_video)
115-
# except Exception as e:
116-
# print(f"Ошибка при объединении видеофайлов: {e}")
117-
#
118-
#
119-
# def add_audio(final_video_path, audio_file):
120-
# """Наложение звука на финальное видео."""
121-
# try:
122-
# final_with_audio = final_video_path.replace(".mp4", "_with_audio.mp4")
123-
# ffmpeg.input(final_video_path).input(audio_file).output(
124-
# final_with_audio, vcodec="copy", acodec="aac", strict="experimental"
125-
# ).run()
126-
# print(f"Финальное видео с наложением звука сохранено в {final_with_audio}")
127-
# except Exception as e:
128-
# print(f"Ошибка при наложении звука: {e}")
129-
#
130-
#
131-
# def main():
132-
# """Основной процесс обработки видео."""
133-
# counter_of_video_batches = START_BATCH_TO_UPSCALE // STEP_PER_BATCH + 1
134-
# start_batch_to_upscale = START_BATCH_TO_UPSCALE
135-
#
136-
# if END_BATCH_TO_UPSCALE == 0:
137-
# input_batches = os.listdir(INPUT_BATCHES_DIR)
138-
# total_batches = len(input_batches)
139-
# input_batches.sort(key=compare_strings)
140-
# frames_in_last_batch = len(
141-
# glob.glob(f"{INPUT_BATCHES_DIR}/{input_batches[-1]}/*jpg")
142-
# )
143-
# end_batch_to_upscale = (
144-
# FRAMES_PER_BATCH * total_batches - 1 + frames_in_last_batch
145-
# )
146-
# else:
147-
# end_batch_to_upscale = END_BATCH_TO_UPSCALE
148-
#
149-
# for batch_num in range(
150-
# start_batch_to_upscale, end_batch_to_upscale, STEP_PER_BATCH
151-
# ):
152-
# start_batch = batch_num
153-
# end_batch = min(batch_num + STEP_PER_BATCH - 1, end_batch_to_upscale)
154-
#
155-
# # Запуск обработки батчей
156-
# upscale_batches(start_batch, end_batch)
157-
#
158-
# # После улучшения фреймов собираем видео
159-
# if MERGE_ALL_BATCHES_TO_VIDEO:
160-
# batches_to_perform = os.listdir(OUTPUT_BATCHES_DIR)
161-
# batches_to_perform.sort(key=compare_strings)
162-
# else:
163-
# batches_to_perform = [f"batch_{i}" for i in range(start_batch, end_batch)]
164-
# batch_video = merge_frames_to_video(
165-
# batches_to_perform, counter_of_video_batches
166-
# )
167-
# if batch_video:
168-
# delete_upscaled_frames()
169-
# combine_videos(TMP_FINAL_VIDEO, batch_video)
170-
# counter_of_video_batches += 1
171-
# print(
172-
# f"Партия фреймов {batch_num // 4 + 1} успешно обработана и объединена."
173-
# )
174-
#
175-
# if os.path.isfile(ORIGINAL_AUDIO):
176-
# add_audio(TMP_FINAL_VIDEO, ORIGINAL_AUDIO)
177-
# else:
178-
# print("Звуковой файл не найден, финальное видео сохранено без звука.")
179-
#
180-
# print(f"Обработка завершена. Итоговое видео сохранено в {TMP_FINAL_VIDEO}")
23+
24+
def merge_frames_to_video(batches_list, video_batch_num, fps):
25+
"""Собирает обработанные фреймы в видео."""
26+
try:
27+
# собираем все апскейленные фреймы из нескольких батчей в 1 список
28+
batches_and_frames_list = list()
29+
for batch in batches_list:
30+
batch_path = f"{OUTPUT_BATCHES_DIR}/{batch}"
31+
batches_and_frames_list.append(
32+
sorted(glob.glob(os.path.join(batch_path, "frame*.jpg")))
33+
)
34+
frames_list = sum(batches_and_frames_list, [])
35+
36+
# файл чтобы скормить пути фреймов ffmpeg
37+
frames_file_path = os.path.join("./", "frames_list.txt")
38+
with open(frames_file_path, "w") as f:
39+
for frame in frames_list:
40+
f.write(f"file '{frame}'\n")
41+
42+
batch_video = os.path.join(
43+
BATCH_VIDEO_PATH, f"batch_{video_batch_num}.mp4"
44+
)
45+
ffmpeg.input(frames_file_path, format="concat", safe=0).output(
46+
batch_video, vcodec="libx264", pix_fmt="yuv420p", crf=18, r=fps
47+
).run()
48+
os.remove(frames_file_path)
49+
print(f"Партия {batches_list} собрана в видео.")
50+
return batch_video
51+
except Exception as e:
52+
print(f"Ошибка сборки видео из фреймов: {e}")
53+
return None
54+
55+
56+
def combine_videos(tmp_final_video, batch_video, fps):
57+
"""Объединяет партии видео в одно целое."""
58+
try:
59+
if os.path.isfile(tmp_final_video):
60+
first_clip = VideoFileClip(tmp_final_video)
61+
second_clip = VideoFileClip(batch_video)
62+
combined_clip = concatenate_videoclips([first_clip, second_clip])
63+
combined_clip.write_videofile(FINAL_VIDEO, codec="libx264", fps=fps)
64+
print(f"Объединение видеофайлов завершено.")
65+
66+
first_clip.close()
67+
second_clip.close()
68+
combined_clip.close()
69+
os.unlink(batch_video)
70+
os.unlink(tmp_final_video)
71+
os.rename(FINAL_VIDEO, tmp_final_video)
72+
else:
73+
shutil.move(batch_video, tmp_final_video)
74+
except Exception as e:
75+
print(f"Ошибка при объединении видеофайлов: {e}")
76+
77+
78+
def main():
79+
"""Основной процесс обработки видео."""
80+
fps = get_fps_accurate(ORIGINAL_VIDEO)
81+
audio = extract_audio()
82+
extract_frames_to_batches()
83+
84+
start_batch = START_BATCH_TO_UPSCALE
85+
end_batch = 0
86+
87+
if END_BATCH_TO_UPSCALE == 0:
88+
# -1 потому что файл .gitkeep в INPUT_BATCHES_DIR тоже считается
89+
end_batch_to_upscale = len(os.listdir(INPUT_BATCHES_DIR)) - 1
90+
else:
91+
end_batch_to_upscale = END_BATCH_TO_UPSCALE
92+
93+
while end_batch != end_batch_to_upscale:
94+
if start_batch + STEP_PER_BATCH <= end_batch_to_upscale:
95+
# если шаг STEP_PER_BATCH доступен в диапазоне, то берем его
96+
end_batch = start_batch + STEP_PER_BATCH - 1
97+
else:
98+
# иначе берем все оставшиеся батчи до конца
99+
end_batch = end_batch_to_upscale
100+
101+
# Запуск обработки батчей
102+
asyncio.run(upscale_batches(start_batch, end_batch))
103+
104+
# После улучшения фреймов собираем видео
105+
# if MERGE_ALL_BATCHES_TO_VIDEO:
106+
# batches_to_perform = os.listdir(OUTPUT_BATCHES_DIR)
107+
# batches_to_perform.sort(key=compare_strings)
108+
# else:
109+
# batches_to_perform = [f"batch_{i}" for i in range(start_batch, end_batch)]
110+
# batch_video = merge_frames_to_video(
111+
# batches_to_perform, counter_of_video_batches
112+
# )
113+
# if batch_video:
114+
# delete_upscaled_frames()
115+
# combine_videos(TMP_FINAL_VIDEO, batch_video)
116+
# counter_of_video_batches += 1
117+
# print(
118+
# f"Партия фреймов {batch_num // 4 + 1} успешно обработана и объединена."
119+
# )
120+
121+
start_batch += STEP_PER_BATCH
122+
123+
if audio:
124+
insert_audio(audio, fps)
125+
else:
126+
print("Звуковой файл не найден, финальное видео сохранено без звука.")
127+
128+
print(f"Обработка завершена. Итоговое видео сохранено в {TMP_FINAL_VIDEO}")
181129

182130

183131
if __name__ == "__main__":
184-
# main()
185-
# fps = get_fps_accurate(ORIGINAL_VIDEO)
186-
# audio = extract_audio()
187-
# insert_audio(audio, fps)
188-
# extract_frames_to_batches()
189-
asyncio.run(upscale_batches(1, 13))
132+
main()

‎src/config/settings.py

+6-1
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,11 @@
3131
ALLOWED_THREADS = int(os.getenv("ALLOWED_THREADS", 6))
3232

3333
# Настройка апскейла
34-
UPSCALE_MODEL_NAME = os.getenv("UPSCALE_MODEL_NAME", f"{ROOT_DIR}/src/utils/realesrgan/models/realesr-animevideov3")
34+
MODEL_DIR = os.getenv("MODEL_DIR", f"{ROOT_DIR}/src/utils/realesrgan/models")
35+
MODEL_NAME = os.getenv("MODEL_NAME", "realesr-animevideov3")
36+
REALESRGAN_SCRIPT = os.getenv(
37+
"REALSESRGAN_DIR",
38+
f"{ROOT_DIR}/src/utils/realesrgan/realesrgan-linux/realesrgan-ncnn-vulkan",
39+
)
3540
UPSCALE_FACTOR = int(os.getenv("UPSCALE_FACTOR", 2))
3641
OUTPUT_IMAGE_FORMAT = os.getenv("OUTPUT_IMAGE_FORMAT", "jpg")

‎src/video_processing/frames.py

+23
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import glob
12
import os
23
import sys
34
from concurrent.futures import ThreadPoolExecutor
@@ -12,6 +13,7 @@
1213
FRAMES_PER_BATCH,
1314
INPUT_BATCHES_DIR,
1415
ORIGINAL_VIDEO,
16+
OUTPUT_IMAGE_FORMAT,
1517
)
1618
from src.utils.batch_utils import make_default_batch_dir
1719

@@ -124,3 +126,24 @@ def extract_frames_to_batches(
124126

125127
video_capture.release()
126128
print("Извлечение завершено.")
129+
130+
131+
def count_total_frames(directory: str) -> int:
132+
"""Считает общее количество фреймов во всех батчах."""
133+
return sum(
134+
len(glob.glob(os.path.join(batch, f"*.{OUTPUT_IMAGE_FORMAT}")))
135+
for batch in glob.glob(os.path.join(directory, "batch_*"))
136+
)
137+
138+
139+
def count_frames_in_certain_batches(directory: str, batches_num_range: range) -> int:
140+
"""Считает общее количество фреймов в указанных батчах."""
141+
frames_in_batches = 0
142+
143+
for batch_num in batches_num_range:
144+
batch_dir = os.path.join(directory, f"batch_{batch_num}")
145+
frames_in_batches += len(
146+
glob.glob(os.path.join(batch_dir, f"*.{OUTPUT_IMAGE_FORMAT}"))
147+
)
148+
149+
return frames_in_batches

0 commit comments

Comments
 (0)