Update generator_function/video_hashid_function.py

generator_function/video_hashid_function.py

@@ -73,7 +73,9 @@ def augment_video_random(
     output_path: Optional[str] = None,
     crf: int = 20,
     preset: str = "medium",
-    variation_idx: int = 0
+    variation_idx: int = 0,
+    border_width: int = 0,
+    border_color: str = "black",
 ) -> str:
     import random
     k_b = random.randint(3, 3)
@@ -85,6 +87,11 @@ def augment_video_random(
 
     vf = f"hue=h={hue_rad:.6f}:s={sat_scale:.4f},eq=contrast={contrast:.4f}:brightness={brightness:.4f}"
 
+    if border_width > 0:
+        ffmpeg_color = border_color.replace("#", "0x") if border_color.startswith("#") else border_color
+        bw = int(border_width)
+        vf += f",pad=iw+{bw*2}:ih+{bw*2}:{bw}:{bw}:{ffmpeg_color}"
+
     inp = Path(input_path)
     suffix = f"_augmented_{variation_idx}" if variation_idx > 0 else "_augmented"
     out = Path(output_path) / f"{inp.stem}{suffix}.mp4" if output_path else inp.with_name(f"{inp.stem}{suffix}.mp4")
@@ -121,12 +128,24 @@ def upload_video_to_r2(video_bytes: bytes, file_name: str) -> str:
     return f"{_public_base_url()}/{key}"
 
 # ---------- PROCESS + UPLOAD ----------
-def process_video_with_hash_info(input_path: str, output_path: Optional[str] = None, variation_idx: int = 0) -> dict:
+def process_video_with_hash_info(
+    input_path: str,
+    output_path: Optional[str] = None,
+    variation_idx: int = 0,
+    border_width: int = 0,
+    border_color: str = "black",
+) -> dict:
     in_sha = _sha256(Path(input_path))
     in_id = _hashid_short(in_sha)
 
     # Augment and get output
-    out_path = augment_video_random(input_path=input_path, output_path=output_path, variation_idx=variation_idx)
+    out_path = augment_video_random(
+        input_path=input_path,
+        output_path=output_path,
+        variation_idx=variation_idx,
+        border_width=border_width,
+        border_color=border_color,
+    )
 
     # output hashid
     out_sha = _sha256(Path(out_path))