From b7037ccccf4e7b4a7e227e8f3a19f92ffc5c152f Mon Sep 17 00:00:00 2001
From: sakamoto <sakamoto@sakamoto.dev>
Date: Sun, 11 Aug 2024 13:31:15 -0500
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20add=20paper=20size?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 utils.py     |  2 ++
 vid2sheet.py | 46 ++++++++++++++++++++++++++++++++++------------
 2 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/utils.py b/utils.py
index 7a2d8ee..ebc2f8f 100644
--- a/utils.py
+++ b/utils.py
@@ -21,6 +21,7 @@ parser.add_argument("-i","--ignore-temp", action="store_true", help="[requires -
 parser.add_argument(
     "-f", "--img-format", type=str, default="png", help="use output custom img format"
 )
+parser.add_argument("-p","--paper-size", type=str, default="A4", help="set paper size")
 
 args = parser.parse_args()
 src = args.source
@@ -29,6 +30,7 @@ verbose = args.verbose
 alt_temp = args.alt_temp
 ignore_temp = args.ignore_temp
 img_format = args.img_format
+paper_size = args.paper_size
 
 if verbose:
     logging.basicConfig(level=logging.DEBUG, format="[%(levelname)s] - %(message)s")
diff --git a/vid2sheet.py b/vid2sheet.py
index b8245be..d459978 100644
--- a/vid2sheet.py
+++ b/vid2sheet.py
@@ -1,7 +1,7 @@
 from utils import Download
 from utils import Misc
 from utils import log
-from utils import src, dest, alt_temp, img_format, ignore_temp, verbose
+from utils import src, dest, alt_temp, img_format, ignore_temp, verbose, paper_size
 
 from tqdm import tqdm
 import tempfile
@@ -17,8 +17,11 @@ import shutil
 # src: str = "https://www.youtube.com/watch?v=tyloC0e-Tqk"
 # dest: str = "/home/sakamoto/Public/test"
 
+
 class Vid2Sheet:
-    def __init__(self, src, dest, img_format: str, use_tempfile=True, ignore_temp=False):
+    def __init__(
+        self, src, dest, img_format: str, use_tempfile=True, ignore_temp=False
+    ):
         self.src = src
         self.dest = dest
         self.img_format = img_format
@@ -72,7 +75,7 @@ class Vid2Sheet:
     def run(self):
         self.check()
         self.capture()
-        self.stitch()
+        self.stitch(paper_size)
         self.convert()
 
     def check(self):
@@ -128,7 +131,7 @@ class Vid2Sheet:
 
         while True:
             ret, current_frame = self.video.read()
-    
+
             if not ret:
                 break
 
@@ -159,10 +162,25 @@ class Vid2Sheet:
         self.video.release()
         log.info("Analysis complete")
 
-    def stitch(self, dpi=300):
+    def stitch(self, paper_size="A4", dpi=300):
         log.info("Attempting to stitch by three for every group...")
-        letter_width = int(8.5 * dpi)
-        letter_height = int(11 * dpi)
+        log.info(f"Chosen paper size: {paper_size}")
+
+        paper_sizes = {
+            "letter": (8.5, 11),
+            "A4": (8.27, 11.69),
+            "legal": (8.5, 14),
+            "tabloid": (11, 17),
+        }
+
+        if paper_size in paper_sizes:
+            width_inches, height_inches = paper_sizes[paper_size]
+        else:
+            log.error(f"Unsupported paper size '{paper_size}'. Using default 'A4'.")
+            width_inches, height_inches = paper_sizes["A4"]
+
+        letter_width = int(width_inches * dpi)
+        letter_height = int(height_inches * dpi)
 
         image_files = sorted(os.listdir(self.raw_dir))
 
@@ -185,18 +203,22 @@ class Vid2Sheet:
                     continue
 
                 img_height, img_width = img.shape[:2]
+
                 scale_factor = min(
                     letter_width / img_width, available_height / img_height
                 )
 
-                resized_img_width = int(img_width * scale_factor)
-                resized_img_height = int(img_height * scale_factor)
+                resized_img_width = int(
+                    img_width * scale_factor * 0.95
+                )
+                resized_img_height = int(img_height * scale_factor * 0.95)
                 resized_img = cv2.resize(img, (resized_img_width, resized_img_height))
 
                 x_offset = (letter_width - resized_img_width) // 2
-                y_offset = (
-                    available_height - resized_img_height
-                ) // 2 + i * available_height
+                y_offset = int(
+                    (available_height - resized_img_height) // 2
+                    + i * available_height * 0.9
+                )
 
                 canvas[
                     y_offset : y_offset + resized_img_height,