Spaces:
Sleeping
Sleeping
| import base64 | |
| import os | |
| import tempfile | |
| # Use temporary directory or current working directory for app data | |
| # This avoids permission issues in containerized environments | |
| # Check if running in container with pre-created directories | |
| if os.path.exists('/tmp/app_data') and os.access('/tmp/app_data', os.W_OK): | |
| app_data_path = '/tmp/app_data' | |
| print("Using pre-created /tmp/app_data directory") | |
| else: | |
| try: | |
| # Try to use /tmp first (usually writable in containers) | |
| app_data_path = '/tmp/app_data' | |
| os.makedirs(app_data_path, exist_ok=True) | |
| except PermissionError: | |
| # Fallback to current directory or temp directory | |
| app_data_path = os.path.join(os.getcwd(), 'app_data') | |
| try: | |
| os.makedirs(app_data_path, exist_ok=True) | |
| except PermissionError: | |
| # Last resort: use system temp directory | |
| app_data_path = tempfile.mkdtemp(prefix='app_data_') | |
| # Make sure the necessary directories exist and have proper permissions | |
| import cv2 | |
| import numpy as np | |
| from flask import Flask, request, jsonify | |
| from mtcnn.mtcnn import MTCNN | |
| from keras_facenet import FaceNet | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| from flask_cors import CORS | |
| from pymongo import MongoClient | |
| from pymongo.server_api import ServerApi | |
| # Initialize MongoDB connection | |
| client = MongoClient('mongodb+srv://nanduvinay719:76qqKRX4zC97yQun@travis.744fuyn.mongodb.net/?retryWrites=true&w=majority&appName=travis', server_api=ServerApi('1')) | |
| if client: | |
| print("Connected to MongoDB") | |
| db = client["travis"] | |
| mongo = db["travis_face_data"] | |
| if "travis_face_data" not in db.list_collection_names(): | |
| db.create_collection("travis_face_data") | |
| app = Flask(__name__) | |
| CORS(app) | |
| # Initialize MTCNN detector and FaceNet model | |
| detector = MTCNN() | |
| embedder = FaceNet() | |
| def cosine(embedding1, embedding2): | |
| dot_product = np.dot(embedding1, embedding2) | |
| norm1 = np.linalg.norm(embedding1) | |
| norm2 = np.linalg.norm(embedding2) | |
| similarity = dot_product / (norm1 * norm2) | |
| return similarity | |
| def reload_embeddings(): | |
| global face_data, labels, names | |
| face_data, labels, names = load_embeddings_from_db() | |
| def home(): | |
| return {"message": "Travis Login API is running!"} | |
| def recognizeLogin(): | |
| try: | |
| if 'image' not in request.files: | |
| return jsonify({"error": "No image provided"}), 400 | |
| file = request.files['image'] | |
| if not file: | |
| return jsonify({"error": "Invalid file"}), 400 | |
| image_data = file.read() | |
| image_array = np.frombuffer(image_data, np.uint8) | |
| image = cv2.imdecode(image_array, cv2.IMREAD_COLOR) | |
| if image is None: | |
| return jsonify({"error": "Invalid image"}), 400 | |
| results = recognize_faces_in_image(image) | |
| print(results[0]) | |
| if results[0]['name'] != 'unknown': | |
| return jsonify({"name": results[0]['name'], "probability": results[0]['probability']}), 200 | |
| else: | |
| return jsonify({'name': "user not recognised"}), 401 | |
| except Exception as e: | |
| print(f"Error in login: {str(e)}") | |
| return jsonify({"error": "Internal server error"}), 500 | |
| def register(): | |
| username = request.form['username'] | |
| phoneno = request.form["phoneNumber"] | |
| email = request.form['email'] | |
| facenet_embeddings = [] | |
| stored_image = None # To store the first grayscale image | |
| print(username) | |
| # Check if user already exists | |
| existing_user = mongo.find_one({"username": username}) | |
| if existing_user: | |
| return jsonify({"error": f"User '{username}' already exists"}), 400 | |
| # Process uploaded images | |
| for i in range(5): # Expecting 5 images | |
| try: | |
| image_file = request.files[f'image{i}'] | |
| except KeyError: | |
| return jsonify({"error": f"Missing image{i} in the request"}), 400 | |
| image_data = image_file.read() | |
| image_array = np.frombuffer(image_data, np.uint8) | |
| image = cv2.imdecode(image_array, cv2.IMREAD_COLOR) | |
| # Face detection using MTCNN for FaceNet | |
| mtcnn_faces = detector.detect_faces(image) | |
| if mtcnn_faces: | |
| # Get the first detected face for FaceNet embedding | |
| x, y, w, h = mtcnn_faces[0]['box'] | |
| x, y = max(0, x), max(0, y) | |
| cropped_face = cv2.resize(image[y:y+h, x:x+w], (160, 160)) | |
| rgb_face = cv2.cvtColor(cropped_face, cv2.COLOR_BGR2RGB) | |
| # Get FaceNet embedding | |
| facenet_embedding = embedder.embeddings(np.expand_dims(rgb_face, axis=0)).flatten() | |
| facenet_embeddings.append(facenet_embedding) | |
| # Face detection using Haar Cascade for CNN | |
| if stored_image is None: | |
| _, buffer = cv2.imencode('.jpg', cv2.cvtColor(cropped_face, cv2.COLOR_BGR2GRAY)) | |
| stored_image = base64.b64encode(buffer).decode('utf-8') | |
| if not facenet_embeddings: | |
| return jsonify({"error": "No valid faces detected in the uploaded images"}), 400 | |
| # Calculate mean embeddings | |
| mean_facenet_embedding = np.mean(facenet_embeddings, axis=0).astype(float).tolist() | |
| # Save model weights with error handling | |
| # Create user data | |
| id = mongo.count_documents({}) + 1 | |
| user_data = { | |
| 'username': username, | |
| 'phoneNumber': phoneno, | |
| 'email': email, | |
| 'embeddings': mean_facenet_embedding, | |
| 'stored_image': stored_image, | |
| 'role': 'agent', | |
| 'id': id | |
| } | |
| # Insert into MongoDB | |
| mongo.insert_one(user_data) | |
| reload_embeddings() | |
| return jsonify({"message": "User registered successfully!"}), 201 | |
| # Load embeddings from MongoDB for recognition | |
| def load_embeddings_from_db(): | |
| try: | |
| users = list(mongo.find({"role": "agent"})) # Only get agent users | |
| face_data = [] # facenet embeddings | |
| labels = [] # id 1,2,3,.. | |
| names = {} # dict of id and username | |
| for user in users: | |
| if 'embeddings' in user: # Only process users with embeddings | |
| face_data.append(user["embeddings"]) | |
| labels.append(user['id']) | |
| names[user['id']] = user['username'] | |
| print(f"Loaded {len(face_data)} user embeddings from database") | |
| return (face_data, labels, names) if face_data else ([], [], {}) | |
| except Exception as e: | |
| print(f"Error loading embeddings from database: {e}") | |
| return [], [], {} | |
| # Load face embeddings from MongoDB initially | |
| face_data, labels, names = load_embeddings_from_db() | |
| def recognize_faces_in_image(image): | |
| if len(face_data) == 0: | |
| return [{"name": "No registered faces", "probability": 0.0}] | |
| faces = detector.detect_faces(image) | |
| results = [] | |
| for face in faces: | |
| x, y, width, height = face['box'] | |
| cropped_face = cv2.resize(image[y:y+height, x:x+width], (160, 160)) | |
| # Convert cropped face to RGB | |
| rgb_face = cv2.cvtColor(cropped_face, cv2.COLOR_BGR2RGB) | |
| embedding = embedder.embeddings(np.expand_dims(rgb_face, axis=0)).flatten() # Use RGB face here | |
| # Compare with stored embeddings in MongoDB | |
| similarities = cosine_similarity([embedding], face_data) | |
| idx = np.argmax(similarities) | |
| best_match = similarities[0][idx] | |
| 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)}) | |
| else: | |
| results.append({"name": "unknown", "probability": float(best_match)}) | |
| return results | |
| def admin_login(): | |
| data = request.get_json() | |
| username = data.get('username') | |
| password = data.get('password') | |
| if not username or not password: | |
| return jsonify({"error": "Username and password are required"}), 400 | |
| # Find admin user in database | |
| admin = mongo.find_one({ | |
| "username": username, | |
| "role": "admin" | |
| }) | |
| if not admin: | |
| return jsonify({"error": "Invalid credentials"}), 401 | |
| # In a real application, you should hash passwords and compare hashes | |
| # For now, we'll just check if the password matches | |
| if admin.get('password') != password: | |
| return jsonify({"error": "Invalid credentials"}), 401 | |
| return jsonify({ | |
| "success": True, | |
| "username": username, | |
| "role": "admin" | |
| }), 200 | |
| if __name__ == '__main__': | |
| app.run(host='0.0.0.0',port=7860,debug=True) | |