Update app.py

app.py

@@ -10,6 +10,11 @@ from tqdm import tqdm
 from streamlit_webrtc import webrtc_streamer, VideoTransformerBase
 import shutil
 import zipfile
+import image_app
+from utils import app
+from embeddings_app import get_embeddings
+import webcam_app
+from webcam_app import FaceRecognitionTransformer
 
 # Function to extract zip file
 def extract_zip(zip_file_path, extract_dir):
@@ -17,88 +22,6 @@ def extract_zip(zip_file_path, extract_dir):
         zip_ref.extractall(extract_dir)
 
 
-class FaceRecognitionTransformer(VideoTransformerBase):
-    def __init__(self):
-        self.app = FaceAnalysis(name='buffalo_l')
-        self.app.prepare(ctx_id=0, det_size=(640, 640))
-        self.names = None
-        self.embeddings = None
-
-    def _recognize_faces(self, frame):
-        if self.names is None or self.embeddings is None:
-            return frame
-
-        # Perform face analysis on the frame
-        faces = self.app.get(frame)
-
-        # Process each detected face separately
-        for face in faces:
-            # Retrieve the embedding for the detected face
-            detected_embedding = face.normed_embedding
-
-            # Calculate similarity scores with known embeddings
-            scores = np.dot(detected_embedding, np.array(self.embeddings).T)
-            scores = np.clip(scores, 0., 1.)
-
-            # Find the index with the highest score
-            idx = np.argmax(scores)
-            max_score = scores[idx]
-
-            # Check if the maximum score is above a certain threshold (adjust as needed)
-            threshold = 0.7
-            if max_score >= threshold:
-                recognized_name = self.names[idx]
-            else:
-                recognized_name = "Unknown"
-
-            # Draw bounding box around the detected face
-            bbox = face.bbox.astype(int)
-            cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2)
-            # Write recognized name within the bounding box
-            cv2.putText(frame, recognized_name, (bbox[0], bbox[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
-            
-            # Debug print
-            print("Detected face:", recognized_name, "with confidence:", max_score)
-
-        return frame
-
-    def transform(self, frame):
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        frame = self._recognize_faces(frame)
-        return frame
-
-# 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 = []
-
-    # Traverse through each subfolder
-    for root, dirs, files in os.walk(db_dir):
-        for folder in dirs:
-            if folder == ".ipynb_checkpoints":
-                continue
-            img_paths = glob(os.path.join(root, 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)
-
-    if embeddings:
-        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)
-    else:
-        st.warning("No embeddings generated. Please ensure that there are valid images with detected faces.")
-
-
 # Function to delete files and directory
 def delete_files(db_dir):
     shutil.rmtree(db_dir)
@@ -140,7 +63,7 @@ def main():
     # Other tabs can be added similarly
     if choice == "Webcam":
         st.header("WEBCAM")
-        st.subheader("upload names and embeddings file")
+        st.subheader("Upload names and embeddings file")
         uploaded_names = st.file_uploader("Upload names.npy", type="npy")
         uploaded_embeddings = st.file_uploader("Upload embeddings.npy", type="npy")
 
@@ -160,7 +83,21 @@ def main():
                 async_processing=True,
             )
 
+    if choice == "Face Recognition in Image":
+        st.header("Image Recognition")
+        st.subheader("Upload names and embeddings file")
+        upload_names = st.file_uploader("Upload names.npy", type="npy")
+        upload_embeddings = st.file_uploader("Upload embeddings.npy", type="npy")
+        st.subheader("Upload Image")
+        upload_img = st.file_uploader("Upload Image",type=["png","jpg"])
 
+        if upload_img and upload_names and upload_embeddings:
+            names = np.load(uploaded_names)
+            embeddings = np.load(uploaded_embeddings)
+            im_array = np.frombuffer(upload_img.read(),np.uint8)
+            img = cv2.imdecode(im_array,cv2.IMREAD_COLOR)
+            if st.button("Verify Faces"):
+                image_app.recognize_and_display(img,embeddings,names,app)
 
 
 if __name__ == "__main__":