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, extract_faces_from_frames_folder 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 # Keep /tmp directory consistent across reruns if "hf_tmp" not in st.session_state: st.session_state.hf_tmp = tempfile.mkdtemp(dir="/tmp") HF_TMP = st.session_state.hf_tmp # Define output paths input_video_path = os.path.join(HF_TMP, "input.mp4") 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") param_log_path = os.path.join(HF_TMP, "grouping_params.txt") 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("👥 How many people are in the video?", min_value=1, value=2) warm_up_models() identitified = False if uploaded_video: # ✅ read once try: with open(input_video_path, "wb") as f: f.write(uploaded_video.read()) st.success("✅ File saved") # st.video(temp_path) except Exception as e: st.error(f"❌ Failed to write video: {e}") if st.button("🚀 Preprocess video"): progress = st.progress(0) with st.spinner("🧠 Warming up models..."): warm_up_models() progress.progress(5) with st.spinner("📽 Extracting frames..."): extract_frames(input_video_path, frames_folder) progress.progress(10) with st.spinner("🕵️ Detecting faces..."): detect_faces_insightface_from_frames( frames_folder, csv_path, streamlit_progress=progress, progress_range=(10, 70) ) progress.progress(70) with st.spinner("✂️ Extracting face crops..."): extract_faces_from_frames_folder( frames_folder, csv_path, output_faces, streamlit_progress=progress, progress_range=(70, 75) ) progress.progress(75) with st.spinner("👥 Grouping identities..."): identitified = group_faces_by_identity_facenet_fixed( faces_folder=output_faces, output_folder=grouped_faces_dir, num_identities=num_ids, streamlit_progress=progress, progress_range=(75, 95), param_log_path=param_log_path ) if identitified: split_csv_by_identity( original_csv=csv_path, grouped_folder=grouped_faces_dir, output_dir=identity_csv_dir ) st.write("📁 Grouped identities:", os.listdir(grouped_faces_dir)) for subfolder in sorted(os.listdir(grouped_faces_dir)): subdir_path = os.path.join(grouped_faces_dir, subfolder) st.write(f" - {subfolder}: {len(os.listdir(subdir_path))} images") progress.progress(95) 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(100) else: st.error(f"❌ Unable to find {num_ids} distinct people in the video. " f"The video may be too short or too low quality. Please try another.") # 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) st.text(f"📍 Current identity index: {current_index}") # Debug if current_index < total_faces: st.header(f"🧬 Identity {current_index + 1}/{total_faces}") # Show original sample image grouped_faces_dir = st.session_state.grouped_faces_dir identity_folders = sorted([ f for f in os.listdir(grouped_faces_dir) if f.startswith("identity_") ]) if current_index < len(identity_folders): sample_dir = os.path.join(grouped_faces_dir, identity_folders[current_index]) else: st.warning("❌ No identity folder found for this index.") sample_dir = None if sample_dir and os.path.exists(sample_dir): st.write(f"📂 Using sample directory: {sample_dir}") 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 (with unique key) if st.button("🔁 Generate New Identity", key=f"generate_{current_index}"): with st.spinner("🧠 Generating new face..."): new_face = generate_face_image() st.session_state.generated_faces[current_index] = new_face st.rerun() # Confirm and move to next identity (with rerun) if st.button("✅ Confirm & Next Identity", key=f"next_{current_index}"): 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 st.experimental_rerun() # 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")