Update FastAPI.py

FastAPI.py

@@ -39,10 +39,18 @@ class ImgInput(BaseModel):
 class ImgOutput(BaseModel):
     user_id: list
 
-
 class UserSaved(BaseModel):
     status: str
 
+class UserDelete(BaseModel):
+    label: str
+
+class Message(BaseModel):
+    message: str
+
+class CleanPickle(BaseModel):
+    confirm: bool
+
 
 def recognize_face(image_url: HttpUrl) -> ImgOutput:
 
@@ -108,6 +116,8 @@ def add_face(image_url: HttpUrl,user_name : str):
     rgb_img = cv2.cvtColor(new_image, cv2.COLOR_BGR2RGB)
 
     encode = face_recognition.face_encodings(new_image)[0]
+    if len(encode) == 0:
+        return "No face found"
     face_encodings.append(encode)
     labels.append(user_name)
 
@@ -122,6 +132,54 @@ def add_face(image_url: HttpUrl,user_name : str):
     # Upload the pickle file to Firebase Storage
     pkl_blob = storage.child(f"{storage_folder}pkl/face_encodings.pkl")
     pkl_blob.put("face_encodings.pkl")
+    return "User Saved"
+
+
+# Function to delete a user
+def delete_user(user_name: str):
+    # Load the stored face encodings and labels from the pickle file
+    with open("face_encodings.pkl", "rb") as file:
+        data = pickle.load(file)
+        face_encodings = data["encodings"]
+        labels = data["labels"]
+
+    if user_name in labels:
+        index = labels.index(user_name)
+        del labels[index]
+        del face_encodings[index]
+
+        # Save the updated encodings and labels to the pickle file
+        data = {"encodings": face_encodings, "labels": labels}
+        with open("face_encodings.pkl", "wb") as file:
+            pickle.dump(data, file)
+
+        # Upload the updated pickle file to Firebase Storage
+        pkl_blob = storage.child(f"{storage_folder}pkl/face_encodings.pkl")
+        pkl_blob.put("face_encodings.pkl")
+
+        return {"message": f"User '{user_name}' deleted successfully."}
+    else:
+        return {"message": f"User '{user_name}' not found."}
+
+
+def clean_pickle(confirm: bool):
+    if confirm:
+        # Remove the pickle file
+        if os.path.exists("face_encodings.pkl"):
+            os.remove("face_encodings.pkl")
+
+        # Create an empty pickle file
+        with open("face_encodings.pkl", "wb") as file:
+            data = {"encodings": [], "labels": []}
+            pickle.dump(data, file)
+
+        # Upload the empty pickle file to Firebase Storage
+        pkl_blob = storage.child(f"{storage_folder}pkl/face_encodings.pkl")
+        pkl_blob.put("face_encodings.pkl")
+
+        return {"message": "Pickle file cleaned and uploaded successfully."}
+    else:
+        return {"message": "Confirmation required to clean the pickle file."}
 
 
 @app.post('/')
@@ -132,6 +190,17 @@ async def scoring_endpoint(item:ImgInput):
 
 @app.post('/user/')
 async def scoring_endpoint(item:ImgSave):
-    add_face(item.image_url, item.user_name)
-    results = "User Saved"
-    return UserSaved(status=results)
+    results = add_face(item.image_url, item.user_name)
+    return Message(message=results)
+
+
+@app.delete('/user_delete/')
+async def scoring_endpoint(item: UserDelete):
+    result = delete_user(item.label)
+    return Message(message=result["message"])
+
+
+@app.delete('/clean/')
+async def clean_pickle_endpoint(item: CleanPickle):
+    result = clean_pickle(item.confirm)
+    return result