Update app.py

app.py

@@ -1,36 +1,85 @@
 import os
-import shutil
+import cv2
+import numpy as np
 import streamlit as st
-import tempfile
+from insightface.app import FaceAnalysis
+from glob import glob
+from tqdm import tqdm
+import shutil
 import zipfile
-import embeddings
 
-def extract_database(zip_file):
-    # Create a temporary directory to extract the contents of the zip file
-    temp_dir = tempfile.mkdtemp()
-    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
-        zip_ref.extractall(temp_dir)
-    return temp_dir
+# Function to extract zip file
+def extract_zip(zip_file_path, extract_dir):
+    with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
+        zip_ref.extractall(extract_dir)
+
+# Function to get embeddings
+def get_embeddings(db_dir):
+    app = FaceAnalysis(name='buffalo_l')
+    app.prepare(ctx_id=0, det_size=(640, 640),)
+    names = []
+    embeddings = []
+
+    folders = os.listdir(db_dir)
+    for folder in tqdm(folders):
+        if ".ipynb_checkpoints" in folder:
+            continue
+        img_paths = glob(os.path.join(db_dir, folder, '*'))
+        for img_path in img_paths:
+            img = cv2.imread(img_path)
+            if img is None:
+                continue
+            faces = app.get(img)
+            if len(faces) != 1:
+                continue
+            face = faces[0]
+            names.append(folder)
+            embeddings.append(face.normed_embedding)
+
+    embeddings = np.stack(embeddings, axis=0)
+    np.save(os.path.join(db_dir, "embeddings.npy"), embeddings)
+    np.save(os.path.join(db_dir, "names.npy"), names)
+
+# Function to delete files and directory
+def delete_files(db_dir):
+    shutil.rmtree(db_dir)
 
+# Main function
 def main():
-    st.title("FACE ANALYZER")
-    tab_selection = st.sidebar.selectbox("Select Functionality", ["Embeddings", "Image Recognition", "Video Recognition", "Webcam Recognition"])
-    if tab_selection == "Embeddings":
-        st.header("Embeddings")
-        zip_file = st.file_uploader("Upload DataBase(ZIP)", type="zip")
-
-        if zip_file:
-            db_dir = extract_database(zip_file)
-            if st.button("Get Embeddings"):
-                db_dir = embeddings.get_embeddings(db_dir)
-                embeddings_path = os.path.join(db_dir, "embeddings.npy")
-                names_path = os.path.join(db_dir, "names.npy")
-
-                st.write("Download numpy files:")
-                st.download_button(label="Download embeddings.npy", data=open(embeddings_path, "rb"), file_name="embeddings.npy")
-                st.download_button(label="Download names.npy", data=open(names_path, "rb"), file_name="names.npy")
-
-            shutil.rmtree(db_dir, ignore_errors=True)
+    st.title("Face Recognition App")
+    # Tabs
+    tabs = ["Embeddings", "Face Recognition in Image", "Face Recognition in Video", "Face Recognition through Webcam"]
+    choice = st.sidebar.selectbox("Select Option", tabs)
+
+    # Embeddings tab
+    if choice == "Embeddings":
+        st.subheader("Upload a Zip File")
+        uploaded_file = st.file_uploader("Choose a zip file", type="zip")
+
+        if uploaded_file is not None:
+            with open("temp.zip", "wb") as f:
+                f.write(uploaded_file.getbuffer())
+            st.success("File uploaded successfully!")
+            extract_zip("temp.zip", "temp")  # Extract the uploaded zip file to the temp directory
+
+        if st.button("Get Embeddings"):
+            get_embeddings("temp")
+            st.success("Embeddings generated successfully!")
+
+        if st.button("Download Files"):
+            st.markdown("### Download names.npy:")
+            with open("temp/names.npy", "rb") as file:
+                st.download_button("Download names.npy", file.getvalue(), file_name="names.npy", mime="application/octet-stream")
+            
+            st.markdown("### Download embeddings.npy:")
+            with open("temp/embeddings.npy", "rb") as file:
+                st.download_button("Download embeddings.npy", file.getvalue(), file_name="embeddings.npy", mime="application/octet-stream")
+
+        if st.button("Delete Files"):
+            delete_files("temp")
+            st.success("Files deleted successfully!")
+
+    # Other tabs can be added similarly
 
 if __name__ == "__main__":
     main()