mirror of
https://github.com/blw1138/Zordon.git
synced 2025-12-17 08:48:13 +00:00
Brett
4df41a2079
* Add a frame complete notification to BaseWorker and distributed_job_manager.py * Add API to download individual files to API server and ServerProxy * Rename subjob notification API and add download_missing_frames_from_subjob * Subjobs will now notify parent when a frame is complete * Fix missed rename * Add some misc logging * Better error handling * Fix frame download file path issue * Download missing frames at job completion and misc cleanup * Misc cleanup * Code cleanup
114 lines
5.0 KiB
Python
114 lines
5.0 KiB
Python
import logging
|
|
import os
|
|
import subprocess
|
|
import threading
|
|
from pathlib import Path
|
|
|
|
from src.utilities.ffmpeg_helper import generate_thumbnail, save_first_frame
|
|
|
|
logger = logging.getLogger()
|
|
supported_video_formats = ['.mp4', '.mov', '.avi', '.mpg', '.mpeg', '.mxf', '.m4v', 'mkv']
|
|
supported_image_formats = ['.jpg', '.png', '.exr', '.tif']
|
|
|
|
|
|
class PreviewManager:
|
|
|
|
storage_path = None
|
|
_running_jobs = {}
|
|
|
|
@classmethod
|
|
def __generate_job_preview_worker(cls, job, replace_existing=False, max_width=320):
|
|
|
|
# Determine best source file to use for thumbs
|
|
job_file_list = job.file_list()
|
|
source_files = job_file_list if job_file_list else [job.input_path]
|
|
preview_label = "output" if job_file_list else "input"
|
|
|
|
# filter by type
|
|
found_image_files = [f for f in source_files if os.path.splitext(f)[-1].lower() in supported_image_formats]
|
|
found_video_files = [f for f in source_files if os.path.splitext(f)[-1].lower() in supported_video_formats]
|
|
|
|
# check if we even have any valid files to work from
|
|
if source_files and not found_video_files and not found_image_files:
|
|
logger.warning(f"No valid image or video files found in files from job: {job}")
|
|
return
|
|
|
|
os.makedirs(cls.storage_path, exist_ok=True)
|
|
base_path = os.path.join(cls.storage_path, f"{job.id}-{preview_label}-{max_width}")
|
|
preview_video_path = base_path + '.mp4'
|
|
preview_image_path = base_path + '.jpg'
|
|
|
|
if replace_existing:
|
|
for x in [preview_image_path, preview_video_path]:
|
|
try:
|
|
os.remove(x)
|
|
except OSError:
|
|
pass
|
|
|
|
# Generate image previews
|
|
if (found_video_files or found_image_files) and not os.path.exists(preview_image_path):
|
|
try:
|
|
path_of_source = found_image_files[-1] if found_image_files else found_video_files[-1]
|
|
logger.debug(f"Generating image preview for {path_of_source}")
|
|
save_first_frame(source_path=path_of_source, dest_path=preview_image_path, max_width=max_width)
|
|
logger.debug(f"Successfully created image preview for {path_of_source}")
|
|
except Exception as e:
|
|
logger.error(f"Error generating image preview for {job}: {e}")
|
|
|
|
# Generate video previews
|
|
if found_video_files and not os.path.exists(preview_video_path):
|
|
try:
|
|
path_of_source = found_video_files[0]
|
|
logger.debug(f"Generating video preview for {path_of_source}")
|
|
generate_thumbnail(source_path=path_of_source, dest_path=preview_video_path, max_width=max_width)
|
|
logger.debug(f"Successfully created video preview for {path_of_source}")
|
|
except subprocess.CalledProcessError as e:
|
|
logger.error(f"Error generating video preview for {job}: {e}")
|
|
|
|
@classmethod
|
|
def update_previews_for_job(cls, job, replace_existing=False, wait_until_completion=False, timeout=None):
|
|
job_thread = cls._running_jobs.get(job.id)
|
|
if job_thread and job_thread.is_alive():
|
|
logger.debug(f'Preview generation job already running for {job}')
|
|
else:
|
|
job_thread = threading.Thread(target=cls.__generate_job_preview_worker, args=(job, replace_existing,))
|
|
job_thread.start()
|
|
cls._running_jobs[job.id] = job_thread
|
|
|
|
if wait_until_completion:
|
|
job_thread.join(timeout=timeout)
|
|
|
|
@classmethod
|
|
def get_previews_for_job(cls, job):
|
|
|
|
results = {}
|
|
try:
|
|
directory_path = Path(cls.storage_path)
|
|
preview_files_for_job = [f for f in directory_path.iterdir() if f.is_file() and f.name.startswith(job.id)]
|
|
|
|
for preview_filename in preview_files_for_job:
|
|
try:
|
|
pixel_width = str(preview_filename).split('-')[-1]
|
|
preview_label = str(os.path.basename(preview_filename)).split('-')[1]
|
|
extension = os.path.splitext(preview_filename)[-1].lower()
|
|
kind = 'video' if extension in supported_video_formats else \
|
|
'image' if extension in supported_image_formats else 'unknown'
|
|
results[preview_label] = results.get(preview_label, [])
|
|
results[preview_label].append({'filename': str(preview_filename), 'width': pixel_width, 'kind': kind})
|
|
except IndexError: # ignore invalid filenames
|
|
pass
|
|
except FileNotFoundError:
|
|
pass
|
|
return results
|
|
|
|
@classmethod
|
|
def delete_previews_for_job(cls, job):
|
|
all_previews = cls.get_previews_for_job(job)
|
|
flattened_list = [item for sublist in all_previews.values() for item in sublist]
|
|
for preview in flattened_list:
|
|
try:
|
|
logger.debug(f"Removing preview: {preview['filename']}")
|
|
os.remove(preview['filename'])
|
|
except OSError as e:
|
|
logger.error(f"Error removing preview '{preview.get('filename')}': {e}")
|