86 lines
4.0 KiB
Python
86 lines
4.0 KiB
Python
import importlib
|
|
import sys
|
|
import modules
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from types import ModuleType
|
|
from typing import Any, List, Callable
|
|
from tqdm import tqdm
|
|
|
|
FRAME_PROCESSORS_MODULES: List[ModuleType] = []
|
|
FRAME_PROCESSORS_INTERFACE = [
|
|
'pre_check',
|
|
'pre_start',
|
|
'process_frame',
|
|
'process_image',
|
|
'process_video'
|
|
]
|
|
|
|
|
|
def load_frame_processor_module(frame_processor: str) -> Any:
|
|
"""Dynamically import a frame processor module and check its interface."""
|
|
try:
|
|
frame_processor_module = importlib.import_module(f'modules.processors.frame.{frame_processor}')
|
|
for method_name in FRAME_PROCESSORS_INTERFACE:
|
|
if not hasattr(frame_processor_module, method_name):
|
|
print(f"Frame processor {frame_processor} missing method: {method_name}")
|
|
sys.exit()
|
|
except ImportError:
|
|
print(f"Frame processor {frame_processor} not found")
|
|
sys.exit()
|
|
return frame_processor_module
|
|
|
|
|
|
def get_frame_processors_modules(frame_processors: List[str]) -> List[ModuleType]:
|
|
"""Get or load all frame processor modules for the given list."""
|
|
global FRAME_PROCESSORS_MODULES
|
|
|
|
if not FRAME_PROCESSORS_MODULES:
|
|
for frame_processor in frame_processors:
|
|
frame_processor_module = load_frame_processor_module(frame_processor)
|
|
FRAME_PROCESSORS_MODULES.append(frame_processor_module)
|
|
set_frame_processors_modules_from_ui(frame_processors)
|
|
return FRAME_PROCESSORS_MODULES
|
|
|
|
|
|
def set_frame_processors_modules_from_ui(frame_processors: List[str]) -> None:
|
|
"""
|
|
Update FRAME_PROCESSORS_MODULES based on UI state.
|
|
Adds or removes frame processor modules according to the UI toggles in modules.globals.fp_ui.
|
|
"""
|
|
global FRAME_PROCESSORS_MODULES
|
|
current_processor_names = [proc.__name__.split('.')[-1] for proc in FRAME_PROCESSORS_MODULES]
|
|
for frame_processor, state in modules.globals.fp_ui.items():
|
|
if state is True and frame_processor not in current_processor_names:
|
|
try:
|
|
frame_processor_module = load_frame_processor_module(frame_processor)
|
|
FRAME_PROCESSORS_MODULES.append(frame_processor_module)
|
|
except SystemExit:
|
|
print(f"SystemExit: Could not load frame processor '{frame_processor}'.")
|
|
except Exception as e:
|
|
print(f"Error loading frame processor '{frame_processor}': {e}")
|
|
elif state is False and frame_processor in current_processor_names:
|
|
try:
|
|
FRAME_PROCESSORS_MODULES = [proc for proc in FRAME_PROCESSORS_MODULES if proc.__name__.split('.')[-1] != frame_processor]
|
|
except Exception as e:
|
|
print(f"Error removing frame processor '{frame_processor}': {e}")
|
|
|
|
|
|
def multi_process_frame(source_path: str, temp_frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None], progress: Any = None) -> None:
|
|
"""Process frames in parallel using a thread pool."""
|
|
with ThreadPoolExecutor(max_workers=modules.globals.execution_threads) as executor:
|
|
futures = []
|
|
for path in temp_frame_paths:
|
|
future = executor.submit(process_frames, source_path, [path], progress)
|
|
futures.append(future)
|
|
for future in futures:
|
|
future.result()
|
|
|
|
|
|
def process_video(source_path: str, frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None]) -> None:
|
|
"""Process a video by processing all frames with a progress bar."""
|
|
progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
|
|
total = len(frame_paths)
|
|
with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress:
|
|
progress.set_postfix({'execution_providers': modules.globals.execution_providers, 'execution_threads': modules.globals.execution_threads, 'max_memory': modules.globals.max_memory})
|
|
multi_process_frame(source_path, frame_paths, process_frames, progress)
|