Update app.py

app.py

@@ -44,10 +44,6 @@ def process_sheet(sheet):
         processed_lines.append(processed_line)
     return '<br>'.join(processed_lines)
 
-# Load a sample image and learn how to recognize it
-known_image = face_recognition.load_image_file("known_face.jpg")
-known_encoding = face_recognition.face_encodings(known_image)[0]
-
 # Main Function
 def main():
     # Layout Configuration
@@ -74,17 +70,28 @@ def main():
             face_locations = face_recognition.face_locations(rgb_image)
             face_encodings = face_recognition.face_encodings(rgb_image, face_locations)
 
+            # Check if the known face image exists
+            if os.path.isfile("known_face.jpg"):
+                known_image = face_recognition.load_image_file("known_face.jpg")
+                known_encoding = face_recognition.face_encodings(known_image)[0]
+            else:
+                known_encoding = None
+
             # Iterate over detected faces and compare with known face
             for (top, right, bottom, left), face_encoding in zip(face_locations, face_encodings):
-                matches = face_recognition.compare_faces([known_encoding], face_encoding)
-
-                if True in matches:
-                    # If a match is found, draw a green rectangle and label
-                    cv2.rectangle(rgb_image, (left, top), (right, bottom), (0, 255, 0), 2)
-                    cv2.putText(rgb_image, "Known Face", (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+                if known_encoding is not None:
+                    matches = face_recognition.compare_faces([known_encoding], face_encoding)
+
+                    if True in matches:
+                        # If a match is found, draw a green rectangle and label
+                        cv2.rectangle(rgb_image, (left, top), (right, bottom), (0, 255, 0), 2)
+                        cv2.putText(rgb_image, "Known Face", (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+                    else:
+                        # If no match, draw a red rectangle
+                        cv2.rectangle(rgb_image, (left, top), (right, bottom), (0, 0, 255), 2)
                 else:
-                    # If no match, draw a red rectangle
-                    cv2.rectangle(rgb_image, (left, top), (right, bottom), (0, 0, 255), 2)
+                    # If no known face is registered, draw a blue rectangle
+                    cv2.rectangle(rgb_image, (left, top), (right, bottom), (255, 0, 0), 2)
 
             # Convert the RGB image back to BGR format for display
             bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
@@ -96,6 +103,11 @@ def main():
                 st.session_state['captured_images'].append(filename)
                 st.session_state['last_captured'] = time.time()
 
+        if st.button("Register Known Face"):
+            if image is not None:
+                cv2.imwrite("known_face.jpg", np.array(image))
+                st.success("Known face registered successfully!")
+
         sidebar_html = "<div style='display:flex;flex-direction:column;'>"
         for img_file in st.session_state['captured_images']:
             image_base64 = get_image_base64(img_file)