Update generator_function/video_hashid_function.py

generator_function/video_hashid_function.py

@@ -67,9 +67,11 @@ def augment_video_random(
     output_path: Optional[str] = None,
     crf: int = 20,
     preset: str = "medium",
+    variation_idx: int = 0
 ) -> str:
     import random
-    k_b, k_c, k_h, k_s = [random.randint(-5, 5) for _ in range(4)]
+    k_b = random.randint(3, 3)
+    k_c, k_h, k_s = [random.randint(-3, 3) for _ in range(3)]
     brightness = max(-1.0, min(1.0, k_b * 0.05))
     contrast = max(0.0, min(2.0, 1.0 + k_c * 0.05))
     hue_rad = math.radians(k_h * 5.0)
@@ -78,7 +80,8 @@ def augment_video_random(
     vf = f"hue=h={hue_rad:.6f}:s={sat_scale:.4f},eq=contrast={contrast:.4f}:brightness={brightness:.4f}"
 
     inp = Path(input_path)
-    out = Path(output_path) / f"{inp.stem}_augmented.mp4" if output_path else inp.with_name(f"{inp.stem}_augmented.mp4")
+    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")
 
     cmd = [
         "ffmpeg", "-nostdin", "-y",
@@ -112,12 +115,12 @@ def upload_video_to_r2(video_bytes: bytes, file_name: str) -> str:
     return f"{os.getenv('NEW_BASE').rstrip('/')}/{key}"
 
 # ---------- PROCESS + UPLOAD ----------
-def process_video_with_hash_info(input_path: str, output_path: Optional[str] = None) -> dict:
+def process_video_with_hash_info(input_path: str, output_path: Optional[str] = None, variation_idx: int = 0) -> 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)
+    out_path = augment_video_random(input_path=input_path, output_path=output_path, variation_idx=variation_idx)
 
     # output hashid
     out_sha = _sha256(Path(out_path))