From 1f329677f77000f8d21ea5b0a48bf3045ddae174 Mon Sep 17 00:00:00 2001
From: jlw4049 <jlw_4049@hotmail.com>
Date: Fri, 1 Mar 2024 16:29:47 -0500
Subject: [PATCH] feat: v1.1.0 feat: added support for specific frames being
 passed via the --frames arg. This should be a single string separated by :
 like so "1001:1002"

---
 frame_forge.py           | 10 ++++++-
 frame_forge/__init__.py  | 63 +++++++++++++++++++++++++++++++++-------
 frame_forge/cli_utils.py |  7 +++++
 3 files changed, 68 insertions(+), 12 deletions(-)
 create mode 100644 frame_forge/cli_utils.py

diff --git a/frame_forge.py b/frame_forge.py
index 7107b2e..2233213 100644
--- a/frame_forge.py
+++ b/frame_forge.py
@@ -3,10 +3,11 @@ from argparse import ArgumentParser
 from frame_forge import GenerateImages
 from frame_forge.exceptions import FrameForgeError
 from frame_forge.utils import exit_application
+from frame_forge.cli_utils import frame_list
 
 
 program_name = "FrameForge"
-__version__ = "1.0.3"
+__version__ = "1.1.0"
 
 
 if __name__ == "__main__":
@@ -18,6 +19,12 @@ if __name__ == "__main__":
 
     parser.add_argument("--source", type=str, help="Path to source file")
     parser.add_argument("--encode", type=str, help="Path to encode file")
+    parser.add_argument(
+        "--frames",
+        type=frame_list,
+        help="Only use this if you want to specify the "
+        "frames to generate, this disables sync frames",
+    )
     parser.add_argument("--image-dir", type=str, help="Path to base image folder")
     parser.add_argument(
         "--indexer",
@@ -83,6 +90,7 @@ if __name__ == "__main__":
         img_generator = GenerateImages(
             source_file=Path(args.source),
             encode_file=Path(args.encode),
+            frames=args.frames,
             image_dir=image_dir,
             indexer=args.indexer,
             index_directory=args.index_dir,
diff --git a/frame_forge/__init__.py b/frame_forge/__init__.py
index 3b11130..3fb1935 100644
--- a/frame_forge/__init__.py
+++ b/frame_forge/__init__.py
@@ -2,7 +2,6 @@ import re
 import shutil
 from random import choice
 from pathlib import Path
-from typing import Tuple
 from numpy import linspace
 from unidecode import unidecode
 import awsmfunc
@@ -16,6 +15,7 @@ class GenerateImages:
         self,
         source_file: Path,
         encode_file: Path,
+        frames: str,
         image_dir: Path,
         indexer: str,
         index_directory: None | str,
@@ -38,6 +38,7 @@ class GenerateImages:
         self.source_node = None
         self.reference_source_file = None
         self.encode_file = encode_file
+        self.frames = frames
         self.encode_node = None
         self.image_dir = image_dir
         self.indexer = indexer
@@ -97,7 +98,9 @@ class GenerateImages:
 
         self.check_de_interlaced(num_source_frames, num_encode_frames)
 
-        b_frames = self.get_b_frames(num_source_frames)
+        b_frames = None
+        if not self.frames:
+            b_frames = self.get_b_frames(num_source_frames)
 
         screenshot_comparison_dir, screenshot_sync_dir = self.generate_folders()
 
@@ -109,15 +112,22 @@ class GenerateImages:
 
         vs_source_info, vs_encode_info = self.handle_subtitles(selected_sub_style)
 
-        img_job = self.generate_screens(
-            b_frames,
-            vs_source_info,
-            vs_encode_info,
-            screenshot_comparison_dir,
-            screenshot_sync_dir,
-            selected_sub_style_ref,
-            selected_sub_style_sync,
-        )
+        if not self.frames:
+            img_job = self.generate_screens(
+                b_frames,
+                vs_source_info,
+                vs_encode_info,
+                screenshot_comparison_dir,
+                screenshot_sync_dir,
+                selected_sub_style_ref,
+                selected_sub_style_sync,
+            )
+        else:
+            img_job = self.generate_exact_screens(
+                vs_source_info,
+                vs_encode_info,
+                screenshot_comparison_dir,
+            )
 
         return img_job
 
@@ -166,6 +176,37 @@ class GenerateImages:
                 callback=self.screen_gen_callback,
             )
 
+    def generate_exact_screens(
+        self,
+        vs_source_info,
+        vs_encode_info,
+        screenshot_comparison_dir,
+    ) -> str:
+        print("\nGenerating screenshots, please wait", flush=True)
+
+        # generate source images
+        awsmfunc.ScreenGen(
+            vs_source_info,
+            frame_numbers=self.frames,
+            fpng_compression=1,
+            folder=screenshot_comparison_dir,
+            suffix="a_source__%d",
+            callback=self.screen_gen_callback,
+        )
+
+        # generate encode images
+        awsmfunc.ScreenGen(
+            vs_encode_info,
+            frame_numbers=self.frames,
+            fpng_compression=1,
+            folder=screenshot_comparison_dir,
+            suffix="b_encode__%d",
+            callback=self.screen_gen_callback,
+        )
+
+        print("Screen generation completed", flush=True)
+        return str(screenshot_comparison_dir)
+
     def generate_screens(
         self,
         b_frames,
diff --git a/frame_forge/cli_utils.py b/frame_forge/cli_utils.py
new file mode 100644
index 0000000..8b416fe
--- /dev/null
+++ b/frame_forge/cli_utils.py
@@ -0,0 +1,7 @@
+import re
+
+
+def frame_list(frames: str) -> list:
+    if not re.match(r"\d+(?::\d+)*$", frames):
+        raise ValueError("Input must be in the format of int:int i.e. 101:104")
+    return [int(x) for x in frames.split(":")]