Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,159 +1,13 @@
 import streamlit as st
-import os
 import tempfile
-from PIL import Image
 
-# VASR pipeline imports
-from generate import generate_face_image
-from faceswap import face_swap_multiple_identities
-from video_utils import extract_frames, recombine_frames
-from face_detection import detect_faces_insightface_from_frames
-from grouping import group_faces_by_identity_facenet_fixed, split_csv_by_identity
-from load_models import warm_up_models
-from main import fix_faces
-from image_utils import add_watermark_to_frames
-
-# Streamlit config
-st.set_page_config(page_title="VASR", layout="centered")
-st.title("🧬 VASR - Video Anonymisation")
-
-# Create a Hugging Face writable temp directory
 HF_TMP = tempfile.mkdtemp()
 
-# Define output paths
-frames_folder = os.path.join(HF_TMP, "frames")
-csv_path = os.path.join(HF_TMP, "frames_detections.csv")
-output_faces = os.path.join(HF_TMP, "output_faces")
-grouped_faces_dir = os.path.join(HF_TMP, "grouped_faces_facenet")
-identity_csv_dir = os.path.join(HF_TMP, "identity_csvs")
-output_frames = os.path.join(HF_TMP, "output_frames")
-output_video_path = os.path.join(HF_TMP, "anonymised_output.mp4")
-
-# UI elements
 uploaded_video = st.file_uploader("📹 Upload video file", type=["mp4", "mov"])
-num_ids = st.number_input("👥 Number of identities", min_value=1, value=2)
-
 if uploaded_video is not None:
-    #video_bytes = uploaded_video.read()  # ✅ read once
-    #st.write("Uploaded size:", len(video_bytes))
-
     with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4", dir=HF_TMP) as tmp_file:
-        tmp_file.write(video_bytes)
+        tmp_file.write(uploaded_video.read())
         input_video_path = tmp_file.name
 
     st.video(input_video_path)
     st.success("✅ Video uploaded and saved.")
-
-
-    if st.button("🚀 Start"):
-        progress = st.progress(0)
-
-        with st.spinner("🧠 Warming up models..."):
-            warm_up_models()
-        progress.progress(10)
-
-        with st.spinner("📽 Extracting frames..."):
-            extract_frames(input_video_path, frames_folder)
-        progress.progress(30)
-
-
-        with st.spinner("🕵️ Detecting faces..."):
-            detect_faces_insightface_from_frames(
-                frames_folder, 
-                csv_path, 
-                streamlit_progress=progress, 
-                progress_range=(30, 50)
-            )
-        progress.progress(50)
-
-        with st.spinner("👥 Grouping identities..."):
-            group_faces_by_identity_facenet_fixed(csv_path, output_faces, grouped_faces_dir)
-            split_csv_by_identity(csv_path, grouped_faces_dir, identity_csv_dir)
-        progress.progress(60)
-
-        st.session_state.identity_csv_paths = [
-            os.path.join(identity_csv_dir, f) for f in sorted(os.listdir(identity_csv_dir)) if f.endswith(".csv")
-        ]
-        st.session_state.grouped_faces_dir = grouped_faces_dir
-        st.session_state.generated_faces = [None] * num_ids
-        st.session_state.identity_index = 0
-        st.session_state.video_ready = True
-
-        st.success("✅ Ready to generate anonymised faces.")
-        progress.progress(70)
-
-
-
-
-# Identity selection loop
-if "generated_faces" in st.session_state:
-    current_index = st.session_state.identity_index
-    total_faces = len(st.session_state.generated_faces)
-
-    if current_index < total_faces:
-        st.header(f"🧬 Identity {current_index + 1}/{total_faces}")
-
-        # Show original sample image
-        sample_dir = os.path.join(grouped_faces_dir, f"identity_{current_index}")
-        sample_path = None
-
-        if os.path.exists(sample_dir):
-            image_files = sorted([
-                f for f in os.listdir(sample_dir)
-                if f.lower().endswith(('.jpg', '.png'))
-            ])
-            if image_files:
-                sample_path = os.path.join(sample_dir, image_files[0])
-        else:
-            st.warning(f"❌ Folder not found: {sample_dir}")
-
-        if sample_path and os.path.exists(sample_path):
-            st.image(sample_path, caption="🎭 Original Face", width=256)
-        else:
-            st.warning("⚠️ No sample image found for this identity.")
-
-        # Show current generated face if exists
-        if st.session_state.generated_faces[current_index] is not None:
-            st.image(st.session_state.generated_faces[current_index], caption="🧬 Generated Face", width=256)
-
-        # Generate new face button
-        if st.button("🔁 Generate New Identity"):
-            with st.spinner("🧠 Generating new face..."):
-                new_face = generate_face_image()
-            st.session_state.generated_faces[current_index] = new_face
-            st.image(new_face, caption="🧬 Generated Face", width=256)
-
-        # Confirm and move to next identity
-        if st.button("✅ Confirm & Next Identity"):
-            if st.session_state.generated_faces[current_index] is None:
-                st.warning("⚠️ Please generate a face before continuing.")
-            else:
-                st.session_state.identity_index += 1
-
-    # If all identities confirmed, show anonymise option
-    if current_index >= total_faces:
-        st.success("✅ All identities reviewed.")
-        if st.button("🚀 Anonymise Video"):
-
-            swap_progress = st.progress(0)
-            with st.spinner("📑 Fixing CSVs..."):
-                identity_csvs = fix_faces(total_faces, st.session_state.identity_csv_paths)
-            swap_progress.progress(1)
-            with st.spinner("🤖 Swapping faces..."):
-                face_swap_multiple_identities(
-                    frame_folder=frames_folder,
-                    output_folder=output_frames,
-                    generated_images=st.session_state.generated_faces,
-                    identity_csv_paths=identity_csvs,
-                    streamlit_progress=swap_progress,
-                    progress_range=(1, 96)
-                )
-                add_watermark_to_frames(output_frames,output_frames)
-
-            swap_progress.progress(96)
-            with st.spinner("🎞 Recombining video..."):
-                recombine_frames(input_video_path, frames_folder=output_frames, output_video=output_video_path)
-            swap_progress.progress(100)
-
-            st.success("🎉 Anonymised video complete!")
-            st.video(output_video_path, format="video/mp4")