Update app_pages/video_hashid_generation.py

app_pages/video_hashid_generation.py

@@ -1,8 +1,9 @@
 import os
+import random
 import zipfile
 import tempfile
 import streamlit as st
-from generator_function.video_hashid_function import process_video_with_hash_info
+from generator_function.video_hashid_function import process_video_with_hash_info, SAFE_EFFECTS
 
 VIDEO_EXTS = (".mp4", ".avi", ".mov", ".mkv", ".wmv", ".flv", ".webm")
 
@@ -14,140 +15,213 @@ def _is_valid_zip_member(name: str) -> bool:
         return False
     return True
 
+def _pick_effects(pool: list, count: int) -> list:
+    """Return `count` effect names drawn from `pool`, distinct within each shuffled pass."""
+    effects = []
+    remaining = []
+    while len(effects) < count:
+        if not remaining:
+            remaining = list(pool)
+            random.shuffle(remaining)
+        effects.append(remaining.pop())
+    return effects
+
+def _collect_video_files(uploaded_file, extract_dir):
+    skipped = []
+    video_files = []
+    is_zip = uploaded_file.name.lower().endswith(".zip")
+
+    if is_zip:
+        zip_path = os.path.join(extract_dir, uploaded_file.name)
+        with open(zip_path, "wb") as f:
+            f.write(uploaded_file.getvalue())
+
+        with zipfile.ZipFile(zip_path, "r") as z:
+            members = [m for m in z.namelist() if _is_valid_zip_member(m)]
+            z.extractall(extract_dir, members)
+
+        for root, _, files in os.walk(extract_dir):
+            if "__MACOSX" in root.split(os.sep):
+                continue
+            for file in sorted(files):
+                if file.startswith("._") or file.startswith("."):
+                    continue
+                if not file.lower().endswith(VIDEO_EXTS):
+                    continue
+                input_path = os.path.join(root, file)
+                try:
+                    if os.path.getsize(input_path) < 1024 * 20:
+                        skipped.append(f"{input_path} (too small)")
+                        continue
+                except OSError:
+                    skipped.append(f"{input_path} (stat failed)")
+                    continue
+                video_files.append(input_path)
+    else:
+        video_path = os.path.join(extract_dir, uploaded_file.name)
+        with open(video_path, "wb") as f:
+            f.write(uploaded_file.getvalue())
+        if os.path.getsize(video_path) < 1024 * 20:
+            skipped.append(f"{uploaded_file.name} (too small)")
+        else:
+            video_files.append(video_path)
+
+    return video_files, skipped
+
+def _generate_and_render(video_files, skipped, output_dir, num_variations, border_width, border_color, effects_pool=None):
+    if effects_pool:
+        num_variations = len(effects_pool)
+    results = []
+    total_tasks = len(video_files) * num_variations
+    progress_bar = st.progress(0, text="Starting...")
+    completed = 0
+
+    for video_num, input_path in enumerate(video_files, 1):
+        effects_for_video = _pick_effects(effects_pool, num_variations) if effects_pool else [None] * num_variations
+        for i in range(num_variations):
+            effect_name = effects_for_video[i]
+            label = (
+                f"Video {video_num} of {len(video_files)}"
+                + (f" — variation {i + 1}/{num_variations}" if num_variations > 1 else "")
+                + (f" [{effect_name}]" if effect_name else "")
+                + f": {os.path.basename(input_path)}"
+            )
+            progress_bar.progress(completed / total_tasks if total_tasks else 1.0, text=label)
+            try:
+                r = process_video_with_hash_info(
+                    input_path=input_path,
+                    output_path=output_dir,
+                    variation_idx=i + 1,
+                    border_width=border_width,
+                    border_color=border_color,
+                    effect_name=effect_name,
+                )
+                if r:
+                    results.append(r)
+            except Exception as e:
+                skipped.append(f"{os.path.basename(input_path)} ({e})")
+            completed += 1
+
+    progress_bar.progress(1.0, text=f"Done — {len(results)} video(s) processed.")
+
+    if not results:
+        st.error("No valid videos were processed. The file may be too small, corrupted, or (for ZIPs) contain only macOS metadata.")
+        if skipped:
+            with st.expander("Skipped / errors"):
+                st.write("\n".join(skipped))
+        return
+
+    st.success(f"Processed {len(results)} video(s). Skipped {len(skipped)} file(s).")
+
+    if skipped:
+        with st.expander("Skipped / errors"):
+            st.write("\n".join(skipped))
+
+
+    if len(results) == 1:
+        r = results[0]
+        with open(r["output_path"], "rb") as f:
+            st.download_button(
+                "Download Augmented Video (MP4)",
+                data=f.read(),
+                file_name=r["output_name"],
+                mime="video/mp4",
+                use_container_width=True
+            )
+    else:
+        final_zip_path = os.path.join(output_dir, "augmented_videos.zip")
+        with zipfile.ZipFile(final_zip_path, "w", compression=zipfile.ZIP_DEFLATED) as outzip:
+            for r in results:
+                outzip.write(r["output_path"], arcname=r["output_name"])
+
+        with open(final_zip_path, "rb") as f:
+            st.download_button(
+                "Download All Augmented Videos (ZIP)",
+                data=f.read(),
+                file_name="augmented_videos.zip",
+                mime="application/zip",
+                use_container_width=True
+            )
+
 def render_video_hashid(uid: str):
     st.markdown("#### Change video hashid")
 
-    uploaded_file = st.file_uploader("Upload a video (.mp4) or a ZIP file of videos", type=["zip", "mp4"])
+    uploaded_file = st.file_uploader(
+        "Upload a video (.mp4) or a ZIP file of videos", type=["zip", "mp4"], key="hashid_uploader"
+    )
 
-    num_variations = st.slider("Number of variations per video", min_value=1, max_value=10, value=1, step=1)
+    num_variations = st.slider(
+        "Number of variations per video", min_value=1, max_value=10, value=1, step=1, key="hashid_variations"
+    )
 
     st.markdown("**Border Options**")
-    enable_border = st.checkbox("Add border to videos", value=False)
+    enable_border = st.checkbox("Add border to videos", value=False, key="hashid_enable_border")
     border_width = 0
     border_color = "#000000"
     if enable_border:
         col1, col2 = st.columns(2)
         with col1:
-            border_width = st.slider("Border width (px)", min_value=1, max_value=50, value=10, step=1)
+            border_width = st.slider(
+                "Border width (px)", min_value=1, max_value=50, value=10, step=1, key="hashid_border_width"
+            )
         with col2:
-            border_color = st.color_picker("Border color", value="#000000")
+            border_color = st.color_picker("Border color", value="#000000", key="hashid_border_color")
 
     generate_clicked = st.button(
         "Generate Augmented Videos",
         type="primary",
         width='stretch',
-        disabled=uploaded_file is None
+        disabled=uploaded_file is None,
+        key="hashid_generate"
     )
 
-    if not (generate_clicked and uploaded_file):
-        return
-
-    is_zip = uploaded_file.name.lower().endswith(".zip")
+    if generate_clicked and uploaded_file:
+        with tempfile.TemporaryDirectory() as extract_dir, tempfile.TemporaryDirectory() as output_dir:
+            video_files, skipped = _collect_video_files(uploaded_file, extract_dir)
+            _generate_and_render(video_files, skipped, output_dir, num_variations, border_width, border_color)
 
-    with tempfile.TemporaryDirectory() as extract_dir, tempfile.TemporaryDirectory() as output_dir:
-        results = []
-        skipped = []
-        video_files = []
+    st.divider()
 
-        if is_zip:
-            zip_path = os.path.join(extract_dir, uploaded_file.name)
-            with open(zip_path, "wb") as f:
-                f.write(uploaded_file.getvalue())
+    # ---------------- Apply Video Filters ----------------
+    st.markdown("#### Apply video filters")
 
-            with zipfile.ZipFile(zip_path, "r") as z:
-                members = [m for m in z.namelist() if _is_valid_zip_member(m)]
-                z.extractall(extract_dir, members)
+    filters_uploaded_file = st.file_uploader(
+        "Upload a video (.mp4) or a ZIP file of videos", type=["zip", "mp4"], key="filters_uploader"
+    )
 
-            # Pre-collect valid video files so we know the total
-            for root, _, files in os.walk(extract_dir):
-                if "__MACOSX" in root.split(os.sep):
-                    continue
-                for file in sorted(files):
-                    if file.startswith("._") or file.startswith("."):
-                        continue
-                    if not file.lower().endswith(VIDEO_EXTS):
-                        continue
-                    input_path = os.path.join(root, file)
-                    try:
-                        if os.path.getsize(input_path) < 1024 * 20:
-                            skipped.append(f"{input_path} (too small)")
-                            continue
-                    except OSError:
-                        skipped.append(f"{input_path} (stat failed)")
-                        continue
-                    video_files.append(input_path)
-        else:
-            video_path = os.path.join(extract_dir, uploaded_file.name)
-            with open(video_path, "wb") as f:
-                f.write(uploaded_file.getvalue())
-            if os.path.getsize(video_path) < 1024 * 20:
-                skipped.append(f"{uploaded_file.name} (too small)")
-            else:
-                video_files.append(video_path)
-
-        total_tasks = len(video_files) * num_variations
-        progress_bar = st.progress(0, text="Starting...")
-        completed = 0
-
-        for video_num, input_path in enumerate(video_files, 1):
-            for i in range(num_variations):
-                label = (
-                    f"Video {video_num} of {len(video_files)}"
-                    + (f" — variation {i + 1}/{num_variations}" if num_variations > 1 else "")
-                    + f": {os.path.basename(input_path)}"
-                )
-                progress_bar.progress(completed / total_tasks if total_tasks else 1.0, text=label)
-                try:
-                    r = process_video_with_hash_info(
-                        input_path=input_path,
-                        output_path=output_dir,
-                        variation_idx=i + 1,
-                        border_width=border_width,
-                        border_color=border_color,
-                    )
-                    if r:
-                        results.append(r)
-                except Exception as e:
-                    skipped.append(f"{os.path.basename(input_path)} ({e})")
-                completed += 1
-
-        progress_bar.progress(1.0, text=f"Done — {len(results)} video(s) processed.")
-
-        if not results:
-            st.error("No valid videos were processed. The file may be too small, corrupted, or (for ZIPs) contain only macOS metadata.")
-            if skipped:
-                with st.expander("Skipped / errors"):
-                    st.write("\n".join(skipped))
-            return
-
-        st.success(f"Processed {len(results)} video(s). Skipped {len(skipped)} file(s).")
+    st.markdown("**Border Options**")
+    filters_enable_border = st.checkbox("Add border to videos", value=False, key="filters_enable_border")
+    filters_border_width = 0
+    filters_border_color = "#000000"
+    if filters_enable_border:
+        col1, col2 = st.columns(2)
+        with col1:
+            filters_border_width = st.slider(
+                "Border width (px)", min_value=1, max_value=50, value=10, step=1, key="filters_border_width"
+            )
+        with col2:
+            filters_border_color = st.color_picker("Border color", value="#000000", key="filters_border_color")
+
+    st.markdown("**Filters**")
+    selected_filters = st.multiselect(
+        "Select filters to apply — one output video is generated per selected filter, per input video",
+        options=list(SAFE_EFFECTS.keys()),
+        default=list(SAFE_EFFECTS.keys()),
+        key="filters_pool"
+    )
 
-        if skipped:
-            with st.expander("Skipped / errors"):
-                st.write("\n".join(skipped))
+    filters_generate_clicked = st.button(
+        "Generate Filtered Videos",
+        type="primary",
+        width='stretch',
+        disabled=filters_uploaded_file is None or not selected_filters,
+        key="filters_generate"
+    )
 
-        if len(results) == 1:
-            r = results[0]
-            with open(r["output_path"], "rb") as f:
-                st.download_button(
-                    "Download Augmented Video (MP4)",
-                    data=f.read(),
-                    file_name=r["output_name"],
-                    mime="video/mp4",
-                    use_container_width=True
-                )
-        else:
-            final_zip_path = os.path.join(output_dir, "augmented_videos.zip")
-            import zipfile as zf
-            with zf.ZipFile(final_zip_path, "w", compression=zf.ZIP_DEFLATED) as outzip:
-                for r in results:
-                    outzip.write(r["output_path"], arcname=r["output_name"])
-
-            with open(final_zip_path, "rb") as f:
-                st.download_button(
-                    "Download All Augmented Videos (ZIP)",
-                    data=f.read(),
-                    file_name="augmented_videos.zip",
-                    mime="application/zip",
-                    use_container_width=True
-                )
+    if filters_generate_clicked and filters_uploaded_file:
+        with tempfile.TemporaryDirectory() as extract_dir, tempfile.TemporaryDirectory() as output_dir:
+            video_files, skipped = _collect_video_files(filters_uploaded_file, extract_dir)
+            _generate_and_render(
+                video_files, skipped, output_dir, len(selected_filters),
+                filters_border_width, filters_border_color, effects_pool=selected_filters
+            )