Compare commits

..

3 Commits

Author SHA1 Message Date
jlw4049
84908723d9 feat: v1.3.0
feat: add arg `--img-lib`. This allows the user to select one of the two image libs, defaulting to fpng.
refactor: improved error handling
feat: added temp directory clean up even if the program crashes
2024-06-20 12:14:06 -04:00
jlw4049
d4f0c46ae3 refactor: re worked how image output folders worked in order to prevent errors from dependency libraries not being able to handle unicode paths in some cases
refactor: removed un-needed comments
feat: add support to select image generation library from the CLI
refactor: only import what we need from awsmfunc instead of the whole module
2024-06-20 12:12:59 -04:00
jlw4049
af5695c647 feat: upgrade vapoursynth from v64 to v65 2024-06-20 12:10:51 -04:00
4 changed files with 121 additions and 58 deletions

@ -7,7 +7,7 @@ from frame_forge.cli_utils import frame_list
program_name = "FrameForge"
__version__ = "1.2.2"
__version__ = "1.3.0"
if __name__ == "__main__":
@ -33,6 +33,13 @@ if __name__ == "__main__":
default="lsmash",
help="Indexer choice",
)
parser.add_argument(
"--img-lib",
type=str,
choices=["imwri", "fpng"],
default="fpng",
help="Image library to use",
)
parser.add_argument(
"--source-index-path", type=str, help="Path to look/create indexes for source"
)
@ -114,6 +121,7 @@ if __name__ == "__main__":
frames=args.frames,
image_dir=image_dir,
indexer=args.indexer,
img_lib=args.img_lib,
source_index_path=args.source_index_path,
encode_index_path=args.encode_index_path,
sub_size=args.sub_size,
@ -134,8 +142,18 @@ if __name__ == "__main__":
subtitle_color=args.subtitle_color,
release_sub_title=args.release_sub_title,
)
except Exception as init_error:
exit_application(f"Initiation Error: {init_error}", 1)
try:
img_gen = img_generator.process_images()
if img_gen:
exit_application(f"Output: {img_gen}", 0)
except FrameForgeError as error:
exit_application(error, 1)
exit_application(f"\nOutput: {img_gen}", 0)
except FrameForgeError as ff_error:
img_generator.clean_temp(False)
exit_application(str(ff_error), 1)
except Exception as except_error:
img_generator.clean_temp(False)
exit_application(f"Unhandled Exception: {except_error}", 1)
finally:
img_generator.clean_temp(False)

@ -1,11 +1,13 @@
import re
import shutil
import tempfile
from random import choice
from pathlib import Path
from typing import Tuple
from numpy import linspace
import awsmfunc
import vapoursynth as vs
from awsmfunc import ScreenGenEncoder, ScreenGen, FrameInfo, DynamicTonemap
from frame_forge.exceptions import FrameForgeError
from frame_forge.utils import get_working_dir, hex_to_bgr
@ -18,6 +20,7 @@ class GenerateImages:
frames: str,
image_dir: Path,
indexer: str,
img_lib: str,
source_index_path: None | str,
encode_index_path: None | str,
sub_size: int,
@ -44,6 +47,7 @@ class GenerateImages:
self.encode_node = None
self.image_dir = image_dir
self.indexer = indexer
self.img_lib = ScreenGenEncoder(img_lib)
self.source_index_path = source_index_path
self.encode_index_path = encode_index_path
self.sub_size = sub_size
@ -65,7 +69,9 @@ class GenerateImages:
self.core = vs.core
self.load_plugins()
def process_images(self):
self.temp_dir: Path = None
def process_images(self) -> Path:
self.check_index_paths()
if self.indexer == "lsmash":
@ -108,7 +114,11 @@ class GenerateImages:
if not self.frames:
b_frames = self.get_b_frames(num_source_frames)
screenshot_comparison_dir, screenshot_sync_dir = self.generate_folders()
(
temp_screenshot_comparison_dir,
temp_selected_dir,
temp_screenshot_sync_dir,
) = self.generate_temp_folders()
self.handle_crop()
@ -119,23 +129,27 @@ class GenerateImages:
vs_source_info, vs_encode_info = self.handle_subtitles(selected_sub_style)
if not self.frames:
img_job = self.generate_screens(
self.generate_screens(
b_frames,
vs_source_info,
vs_encode_info,
screenshot_comparison_dir,
screenshot_sync_dir,
temp_screenshot_comparison_dir,
temp_screenshot_sync_dir,
selected_sub_style_ref,
selected_sub_style_sync,
)
else:
img_job = self.generate_exact_screens(
self.generate_exact_screens(
vs_source_info,
vs_encode_info,
screenshot_comparison_dir,
temp_screenshot_comparison_dir,
)
return img_job
final_folder = self.generate_final_folder()
self.move_images(temp_screenshot_comparison_dir.parent, final_folder)
self.clean_temp()
return final_folder
@staticmethod
def screen_gen_callback(sg_call_back):
@ -154,13 +168,14 @@ class GenerateImages:
text=f"Reference\nFrame: {ref_frame}",
style=selected_sub_style_ref,
)
awsmfunc.ScreenGen(
ScreenGen(
vs_encode_ref_info,
frame_numbers=[ref_frame],
fpng_compression=1,
folder=screenshot_sync_dir,
suffix="b_encode__%d",
callback=self.screen_gen_callback,
encoder=self.img_lib,
)
def generate_sync_screens(
@ -173,13 +188,14 @@ class GenerateImages:
text=f"Sync\nFrame: {sync_frame}",
style=selected_sub_style_sync,
)
awsmfunc.ScreenGen(
ScreenGen(
vs_sync_info,
frame_numbers=[sync_frame],
fpng_compression=1,
folder=Path(screenshot_sync_dir),
suffix="a_source__%d",
callback=self.screen_gen_callback,
encoder=self.img_lib,
)
def generate_exact_screens(
@ -187,31 +203,33 @@ class GenerateImages:
vs_source_info,
vs_encode_info,
screenshot_comparison_dir,
) -> str:
) -> Path:
print("\nGenerating screenshots, please wait", flush=True)
# generate source images
awsmfunc.ScreenGen(
ScreenGen(
vs_source_info,
frame_numbers=self.frames,
fpng_compression=1,
folder=screenshot_comparison_dir,
suffix="a_source__%d",
callback=self.screen_gen_callback,
encoder=self.img_lib,
)
# generate encode images
awsmfunc.ScreenGen(
ScreenGen(
vs_encode_info,
frame_numbers=self.frames,
fpng_compression=1,
folder=screenshot_comparison_dir,
suffix="b_encode__%d",
callback=self.screen_gen_callback,
encoder=self.img_lib,
)
print("Screen generation completed", flush=True)
return str(screenshot_comparison_dir)
return screenshot_comparison_dir
def generate_screens(
self,
@ -222,7 +240,7 @@ class GenerateImages:
screenshot_sync_dir,
selected_sub_style_ref,
selected_sub_style_sync,
) -> str:
) -> Path:
print("\nGenerating screenshots, please wait", flush=True)
# handle re_sync if needed
@ -239,23 +257,25 @@ class GenerateImages:
sync_frames = b_frames
# generate source images
awsmfunc.ScreenGen(
ScreenGen(
vs_source_info,
frame_numbers=sync_frames,
fpng_compression=1,
folder=screenshot_comparison_dir,
suffix="a_source__%d",
callback=self.screen_gen_callback,
encoder=self.img_lib,
)
# generate encode images
awsmfunc.ScreenGen(
ScreenGen(
vs_encode_info,
frame_numbers=b_frames,
fpng_compression=1,
folder=screenshot_comparison_dir,
suffix="b_encode__%d",
callback=self.screen_gen_callback,
encoder=self.img_lib,
)
# generate some sync frames
@ -294,13 +314,13 @@ class GenerateImages:
)
print("Screen generation completed", flush=True)
return str(screenshot_comparison_dir)
return screenshot_comparison_dir
def handle_subtitles(self, selected_sub_style):
vs_source_info = self.core.sub.Subtitle(
clip=self.source_node, text="Source", style=selected_sub_style
)
vs_encode_info = awsmfunc.FrameInfo(
vs_encode_info = FrameInfo(
clip=self.encode_node,
title=self.release_sub_title if self.release_sub_title else "",
style=selected_sub_style,
@ -310,10 +330,8 @@ class GenerateImages:
def handle_hdr(self):
if self.tone_map:
self.source_node = awsmfunc.DynamicTonemap(
clip=self.source_node, libplacebo=False
)
self.encode_node = awsmfunc.DynamicTonemap(
self.source_node = DynamicTonemap(clip=self.source_node, libplacebo=False)
self.encode_node = DynamicTonemap(
clip=self.encode_node,
reference=self.reference_source_file,
libplacebo=False,
@ -370,8 +388,8 @@ class GenerateImages:
bottom=self.bottom_crop if self.bottom_crop else 0,
)
def generate_folders(self):
print("\nCreating folders for images", flush=True)
def generate_final_folder(self) -> Path:
print("\nCreating final output folder", flush=True)
if self.image_dir:
image_output_dir = Path(self.image_dir)
else:
@ -379,31 +397,58 @@ class GenerateImages:
Path(self.encode_file).parent / f"{Path(self.encode_file).stem}_images"
)
# check if temp image dir exists, if so delete it!
if image_output_dir.exists():
shutil.rmtree(image_output_dir, ignore_errors=True)
# create main image dir
image_output_dir.mkdir(exist_ok=True, parents=True)
# create comparison image directory and define it as variable
Path(Path(image_output_dir) / "img_comparison").mkdir(exist_ok=True)
screenshot_comparison_dir = str(Path(Path(image_output_dir) / "img_comparison"))
# create selected image directory and define it as variable
Path(Path(image_output_dir) / "img_selected").mkdir(exist_ok=True)
# create sync image directory and define it as variable
Path(Path(image_output_dir) / "img_sync").mkdir(exist_ok=True)
screenshot_sync_dir = str(Path(Path(image_output_dir) / "img_sync"))
# create sub directories
Path(Path(image_output_dir) / "img_sync/sync1").mkdir(exist_ok=True)
Path(Path(image_output_dir) / "img_sync/sync2").mkdir(exist_ok=True)
print("Folder creation completed", flush=True)
return screenshot_comparison_dir, screenshot_sync_dir
return image_output_dir
def generate_temp_folders(self) -> Tuple[Path, Path, Path]:
print("\nCreating temporary folders for images", flush=True)
self.temp_dir = Path(tempfile.mkdtemp(prefix="ff_"))
screenshot_comparison_dir = Path(Path(self.temp_dir) / "img_comparison")
screenshot_comparison_dir.mkdir(exist_ok=True)
selected_dir = Path(Path(self.temp_dir) / "img_selected")
selected_dir.mkdir(exist_ok=True)
screenshot_sync_dir = Path(Path(self.temp_dir) / "img_sync")
screenshot_sync_dir.mkdir(exist_ok=True)
Path(screenshot_sync_dir / "sync1").mkdir(exist_ok=True)
Path(screenshot_sync_dir / "sync2").mkdir(exist_ok=True)
print("Folder creation completed", flush=True)
return screenshot_comparison_dir, selected_dir, screenshot_sync_dir
def move_images(self, temp_folder: Path, output_folder: Path) -> None:
print("\nMoving generated images")
for sub_folder in temp_folder.iterdir():
if sub_folder.is_dir():
target_sub_folder = output_folder / sub_folder.name
target_sub_folder.mkdir(parents=True, exist_ok=True)
for item in sub_folder.iterdir():
target_item = target_sub_folder / item.name
if item.is_dir():
shutil.move(item, target_item)
else:
shutil.move(item, target_sub_folder)
print("Image move completed", flush=True)
def clean_temp(self, status: bool = True) -> None:
if status:
print("\nRemoving temp folder")
shutil.rmtree(self.temp_dir, ignore_errors=True)
if status:
print("Temp folder removal completed")
def get_b_frames(self, num_source_frames):
print(

16
poetry.lock generated

@ -280,28 +280,28 @@ testing-integration = ["build[virtualenv] (>=1.0.3)", "filelock (>=3.4.0)", "jar
[[package]]
name = "vapoursynth"
version = "64"
version = "65"
description = "A frameserver for the 21st century"
optional = false
python-versions = "*"
files = [
{file = "VapourSynth-64-cp311-cp311-win_amd64.whl", hash = "sha256:ad046a704537276f7ebb4132e1210e8922ec4b2403a4b44f5fe30fd9456a6412"},
{file = "VapourSynth-64-cp38-cp38-win_amd64.whl", hash = "sha256:794f93bcaa1ce79510c11962f677a7d18685a0c29aa36586cb307d1eb9d7f2e0"},
{file = "VapourSynth-64.zip", hash = "sha256:29425a135ca68cbb17b3dfcad0097375b75f94af1f499ba89bcd2b03b2651846"},
{file = "VapourSynth-65-cp311-cp311-win_amd64.whl", hash = "sha256:9fab72bb6dbcdd8d0e6643da68908f85387bc764b836524e97a1e6e8989ca13e"},
{file = "VapourSynth-65-cp38-cp38-win_amd64.whl", hash = "sha256:d9b49b595dc929d63250bd82f05e75238b6dfc4a50ab08e7fc4fe4f8888f6b95"},
{file = "VapourSynth-65.zip", hash = "sha256:1d42d461ef9988a3477134e478a2291d79f3469635cde8af2c66b1e87c36f711"},
]
[[package]]
name = "vapoursynth-portable"
version = "64"
version = "65"
description = "A frameserver for the 21st century"
optional = false
python-versions = "*"
files = [
{file = "VapourSynth_portable-64-py2.py3-none-win_amd64.whl", hash = "sha256:82b26b38774197a7ef53cbd3206bafb02d197100fad513d635bf83f9981dad6c"},
{file = "VapourSynth_portable-65-py2.py3-none-win_amd64.whl", hash = "sha256:6a66615d8a25a71766d035054f9980fd8be47dbd97332e9c068933d1b25061b3"},
]
[package.dependencies]
vapoursynth = "64"
vapoursynth = "65"
[[package]]
name = "vs-rekt"
@ -335,4 +335,4 @@ vapoursynth = "*"
[metadata]
lock-version = "2.0"
python-versions = "3.11.5"
content-hash = "759b5ce3003b78e4f9730f1df864a030c6c964151623a94a87cfa4b4f258c1ed"
content-hash = "4425fd0e4eb9ee21b1b0d79c3824a831cdbe45b6733c9477065c8cdd265731a8"

@ -8,9 +8,9 @@ readme = "README.md"
[tool.poetry.dependencies]
python = "3.11.5"
awsmfunc = "^1.3.4"
vapoursynth-portable = "64"
vapoursynth-portable = "65"
numpy = "^1.26.2"
awsmfunc = "^1.3.4"
[tool.poetry.group.dev.dependencies]