Spaces:
Runtime error
Runtime error
File size: 4,083 Bytes
29572ad | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | 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()
|