Update app.py

app.py

@@ -57,6 +57,11 @@ from datetime import datetime
 app = Flask(__name__)
 CORS(app)
 
+# Role management settings
+DEFAULT_ROLE = "user"
+ALLOWED_ROLES = {"user", "admin"}
+DEFAULT_ADMIN_ROLLNUMBER = os.getenv("DEFAULT_ADMIN_ROLLNUMBER", "23BD1A056D").upper()
+
 # Initialize MTCNN detector and FaceNet model
 detector = MTCNN()
 embedder = FaceNet()
@@ -68,6 +73,22 @@ app.config["MONGO_URI"] = "mongodb+srv://nanduvinay719:76qqKRX4zC97yQun@travis.7
 mongo = PyMongo(app)#initialize
 haar_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
 
+def normalize_role(role_value, rollnumber=None):
+    role = (role_value or "").strip().lower()
+    if role in ALLOWED_ROLES:
+        return role
+    if rollnumber and str(rollnumber).upper() == DEFAULT_ADMIN_ROLLNUMBER:
+        return "admin"
+    return DEFAULT_ROLE
+
+def count_admins():
+    users = mongo.db.data.find({}, {"role": 1, "RollNumber": 1})
+    return sum(
+        1
+        for user in users
+        if normalize_role(user.get("role"), user.get("RollNumber")) == "admin"
+    )
+
 def create_cnn_embedding_model():
     model = Sequential([
         Conv2D(32, (3, 3), activation='relu', input_shape=(160, 160, 1)),
@@ -106,11 +127,15 @@ def cnnlogin():
 
     # Fetch user data from MongoDB
     # Collection name 'data' is already correctly used here
-    user_data = mongo.db.data.find_one({"RollNumber": rollnumber}, {"CNN_embeddings": 1, "username": 1})
+    user_data = mongo.db.data.find_one(
+        {"RollNumber": rollnumber},
+        {"CNN_embeddings": 1, "username": 1, "role": 1, "RollNumber": 1},
+    )
     if user_data is None:
         return jsonify({"error": "User not found"}), 404
     stored_cnn_embedding = np.array(user_data["CNN_embeddings"])# Convert stored embeddings to NumPy array
     username = user_data["username"]
+    role = normalize_role(user_data.get("role"), user_data.get("RollNumber", rollnumber))
     print(username)
     # Decode image
     image_data = file.read()
@@ -145,7 +170,9 @@ def cnnlogin():
     # Set a threshold for recognition
     recognition_threshold = 0.94  # Adjust this threshold as needed
     if similarity > recognition_threshold:
-        return jsonify({"name": username, "probability": float(similarity)}), 200
+        return jsonify(
+            {"name": username, "probability": float(similarity), "role": role}
+        ), 200
     else:
         return jsonify({"error": "Face not recognized", "probability": float(similarity)}), 401
 
@@ -175,13 +202,27 @@ def recognizeLogin():
                 {'$set': {today: True}},  # Mark as present for today
                 upsert=True
             )
-        return jsonify({"name": name, "probability": results[0]['probability']}), 200
+        user_data = mongo.db.data.find_one(
+            {"RollNumber": name},
+            {"role": 1, "RollNumber": 1},
+        )
+        role = normalize_role(
+            user_data.get("role") if user_data else None,
+            user_data.get("RollNumber") if user_data else name,
+        )
+        return jsonify(
+            {
+                "name": name,
+                "probability": results[0]["probability"],
+                "role": role,
+            }
+        ), 200
     else:
         return jsonify({'name':"user not recognised"})
 
 @app.route('/register', methods=['POST'])
 def register():
-    rollnumber = request.form["RollNumber"]
+    rollnumber = request.form["RollNumber"].upper()
     username = request.form['Username']
     fathername = request.form["FatherName"]
     phoneno = request.form["phoneNumber"]
@@ -191,7 +232,10 @@ def register():
     print(username)
     # Check if user already exists
     # Collection name 'data' is already correctly used here
-    existing_user = mongo.db.data.find_one({"username": username}, {"_id": 0})
+    existing_user = mongo.db.data.find_one(
+        {"$or": [{"username": username}, {"RollNumber": rollnumber}]},
+        {"_id": 0},
+    )
     if existing_user:
         return jsonify({"error": f"User '{username}' already exists"}), 400
 
@@ -263,6 +307,7 @@ def register():
         'username': username,
         'FatherName': fathername,
         'phoneNumber': phoneno,
+        'role': DEFAULT_ROLE,
         'embeddings': mean_facenet_embedding,
         'CNN_embeddings': mean_cnn_embedding,
         'stored_image': stored_image,
@@ -282,34 +327,48 @@ def register():
 @app.route('/get_users', methods=['GET'])
 def get_users():
     # Collection name 'data' is already correctly used here
-    users = list(mongo.db.data.find({}, {"id": 1, "username": 1}))  # Fetch only _id and username
+    users = list(
+        mongo.db.data.find({}, {"id": 1, "username": 1, "RollNumber": 1, "role": 1})
+    )
     user_count = len(users)
-    return jsonify({"users": [user['username'] for user in users], "count": user_count})
+    formatted_users = []
+    for user in users:
+        roll = user.get("RollNumber", "")
+        formatted_users.append(
+            {
+                "rollNumber": roll,
+                "username": user.get("username", ""),
+                "role": normalize_role(user.get("role"), roll),
+            }
+        )
+    return jsonify({"users": formatted_users, "count": user_count})
 
 def load_embeddings_from_db():
     if mongo is None:
         print("MongoDB not connected - returning empty embeddings")
-        return [], [], {}
+        return [], [], {}, {}
     
     try:
         users = list(mongo.db.data.find())
         face_data = []# facenet embeddings
         labels = [] # id 1,2,3,..
         names = {} #dict of id and roll number
+        roles = {}
         # {"1":vinay,"2":shahank}
         for user in users:
             face_data.append(user["embeddings"])
             labels.append(user['id'])  # Keep the ObjectId
             names[user['id']] = user['RollNumber']  # Use ObjectId as key
+            roles[user['id']] = normalize_role(user.get("role"), user.get("RollNumber"))
 
         print(f"Loaded {len(face_data)} user embeddings from database")
-        return (face_data, labels, names) if face_data else ([], [], {})
+        return (face_data, labels, names, roles) if face_data else ([], [], {}, {})
     except Exception as e:
         print(f"Error loading embeddings from database: {e}")
-        return [], [], {}
+        return [], [], {}, {}
 
 # Load face embeddings from MongoDB initially
-face_data, labels, names = load_embeddings_from_db()
+face_data, labels, names, roles = load_embeddings_from_db()
 
 # Load CNN model weights with error handling
 try:
@@ -327,7 +386,7 @@ except Exception as e:
 
 # Reload embeddings to update after a new registration
 def reload_embeddings():
-    global face_data, labels, names
+    global face_data, labels, names, roles
     try:
         weights_path = os.path.join(app_data_path, 'cnn_model.weights.h5')
         if os.path.exists(weights_path):
@@ -337,7 +396,7 @@ def reload_embeddings():
     except Exception as e:
         print(f"Warning: Could not reload CNN model weights: {e}")
     
-    face_data, labels, names = load_embeddings_from_db()
+    face_data, labels, names, roles = load_embeddings_from_db()
 
 # Recognize faces using MongoDB-stored embeddings
 model = YOLO('yolov5s.pt')  # Replace with your YOLO model path
@@ -374,7 +433,7 @@ def upload_image():
 
 def recognize_faces_in_image(image):
     if len(face_data) == 0:
-        return [{"name": "No registered faces", "probability": 0.0}]
+        return [{"name": "No registered faces", "probability": 0.0, "role": "unknown"}]
 
     faces = detector.detect_faces(image)
     results = []
@@ -394,11 +453,52 @@ def recognize_faces_in_image(image):
         if best_match > 0.7:
             recognized_id = labels[idx] # Get the ObjectId
             recognized_name = names[recognized_id]  # Use ObjectId to get the username
-            results.append({"name": recognized_name, "probability": float(best_match)})
+            role = roles.get(recognized_id, DEFAULT_ROLE)
+            results.append(
+                {"name": recognized_name, "probability": float(best_match), "role": role}
+            )
         else:
-            results.append({"name": "Unknown", "probability": float(best_match)})
+            results.append(
+                {"name": "Unknown", "probability": float(best_match), "role": "unknown"}
+            )
     return results
 
+@app.route('/users/<rollnumber>/role', methods=['PATCH'])
+def update_user_role(rollnumber):
+    payload = request.get_json(silent=True) or {}
+    role = str(payload.get("role", "")).strip().lower()
+    if role not in ALLOWED_ROLES:
+        return jsonify({"error": "Invalid role"}), 400
+
+    rollnumber = str(rollnumber).upper()
+    user = mongo.db.data.find_one({"RollNumber": rollnumber})
+    if user is None:
+        return jsonify({"error": "User not found"}), 404
+
+    current_role = normalize_role(user.get("role"), user.get("RollNumber"))
+    if current_role == "admin" and role != "admin" and count_admins() <= 1:
+        return jsonify({"error": "Cannot remove the last admin"}), 400
+
+    mongo.db.data.update_one({"_id": user["_id"]}, {"$set": {"role": role}})
+    reload_embeddings()
+    return jsonify({"message": "Role updated", "role": role})
+
+@app.route('/users/<rollnumber>', methods=['DELETE'])
+def delete_user(rollnumber):
+    rollnumber = str(rollnumber).upper()
+    user = mongo.db.data.find_one({"RollNumber": rollnumber})
+    if user is None:
+        return jsonify({"error": "User not found"}), 404
+
+    current_role = normalize_role(user.get("role"), user.get("RollNumber"))
+    if current_role == "admin" and count_admins() <= 1:
+        return jsonify({"error": "Cannot delete the last admin"}), 400
+
+    mongo.db.data.delete_one({"_id": user["_id"]})
+    mongo.db.attendance1.delete_one({"username": rollnumber})
+    reload_embeddings()
+    return jsonify({"message": "User deleted"})
+
 @app.route('/users/<username>/images', methods=['GET'])
 def get_user_images(username):
     username = str(username).upper()
@@ -477,4 +577,4 @@ def get_user_attendance(username):
 def get_attendance():
     # Collection name 'attendance1' is already correctly used here
     records = list(mongo.db.attendance1.find({}, {"_id": 0}))
-    return jsonify({"attendance": records})
+    return jsonify({"attendance": records})