image_capture.py

import warnings
warnings.filterwarnings("ignore")
import os
import sys
import streamlit as st
import cv2
import string
import random
import shutil
from PIL import Image
from pydub import AudioSegment
from streamlit_webrtc import webrtc_streamer, WebRtcMode, RTCConfiguration, VideoProcessorBase
from streamlit_extras.switch_page_button import switch_page

# Add the faceRecognize module path to the system path
sys.path.append(os.path.abspath('src/faceRecognize'))

# Suppress TensorFlow logging messages
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'

def generate_unique_filename():
    chars = string.ascii_letters + string.digits
    random_string = ''.join(random.choices(chars, k=8))
    return random_string + '.jpg'

def load_alert_sound():
    song = AudioSegment.from_file('src/faceRecognize/audio/alert.wav', format="wav")
    return song

def create_new_folder():
    initial_folder_path = "src/faceRecognize/facerec/data"
    new_folder_name = st.text_input("Enter new folder name:")
    folder_path = None
    if new_folder_name and st.button("Create Folder"):
        new_folder_path = os.path.join(initial_folder_path, new_folder_name)
        os.makedirs(new_folder_path, exist_ok=True)
        st.success(f"Folder '{new_folder_name}' created successfully.")
        folder_path = new_folder_path
    return folder_path

def choose_existing_folder():
    initial_folder_path = "src/faceRecognize/facerec/data"
    folder_path = st.text_input("Choose from existing user:", initial_folder_path)
    if folder_path and st.button("Select Folder"):
        st.success(f"You selected the folder: {folder_path}")
    return folder_path

class ImageCaptureProcessor(VideoProcessorBase):
    def __init__(self):
        self.frame = None

    def recv(self, frame):
        self.frame = frame.to_ndarray(format="bgr24")
        return frame

    def get_frame(self):
        return self.frame

def capture_and_save_image(path, processor):
    if processor:
        frame = processor.get_frame()
        if frame is not None:
            st.image(frame, caption="Captured Image", channels="BGR", use_column_width=True)
            if st.button("Save Image"):
                filename = generate_unique_filename()
                filepath = os.path.join(path, filename)
                os.makedirs(path, exist_ok=True)
                cv2.imwrite(filepath, frame)
                st.success(f"Picture saved successfully as {filename}")
        else:
            st.warning("No frame captured. Please try again.")

def delete_folder():
    initial_folder_path = "src/faceRecognize/facerec/data"
    folder_path = st.text_input("Enter folder path to delete:", initial_folder_path)
    if folder_path and st.button("Delete Folder"):
        try:
            shutil.rmtree(folder_path)
            st.success("Folder deleted successfully.")
        except Exception as e:
            st.error(f"Failed to delete folder: {str(e)}")

st.title("Face Recognition App")
# RTC configuration for WebRTC
RTC_CONFIGURATION = RTCConfiguration({
    "iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]
})
def main(path):
    ctx = webrtc_streamer(
        key="snapshot",
        mode=WebRtcMode.SENDRECV,
        rtc_configuration=RTC_CONFIGURATION,
        video_processor_factory=ImageCaptureProcessor,
        media_stream_constraints={"video": True, "audio": False}
    )
    print("Calling Main")
    if ctx.video_processor:
        if st.button("Capture Image"):
            capture_and_save_image(path, ctx.video_processor)

option = st.sidebar.selectbox("Choose an option", 
                    ("Run Code", "Create User/Add Photos", "Delete User", "Start Training"))

if option == "Create User/Add Photos":
    sub_option = st.sidebar.selectbox("Choose an option", ("Choose From Existing", "Add User"))
    path = None
    if sub_option == "Choose From Existing":
        path = choose_existing_folder()
    else:
        path = create_new_folder()
    
    if path:
        main(path)

elif option == "Delete User":
    delete_folder()