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

frame_forge.py

@@ -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)

frame_forge/__init__.py

@@ -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(

poetry.lock

@@ -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"

pyproject.toml

@@ -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]