app.py

from flask import Flask, request, render_template, Response, jsonify, url_for, send_from_directory
import cv2
import numpy as np
import joblib
import os
import warnings
from werkzeug.utils import secure_filename
import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
from mediapipe.framework.formats import landmark_pb2
from flask_cors import CORS
from dotenv import load_dotenv
# Import database functions
from database import upload_image_to_cloudinary, save_analysis_to_db

# Load environment variables
load_dotenv()

# Suppress specific deprecation warnings from protobuf
warnings.filterwarnings("ignore", category=UserWarning, module='google.protobuf')
app = Flask(__name__, template_folder='templates')
CORS(app)  # Enable CORS for all routes

# Configure Flask for production
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max file size
app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'}

# Initialize MediaPipe Face Landmarker (only once)
try:
    base_options = python.BaseOptions(model_asset_path='face_landmarker_v2_with_blendshapes.task')
    options = vision.FaceLandmarkerOptions(base_options=base_options,
                                         output_face_blendshapes=True,
                                         output_facial_transformation_matrixes=True,
                                         num_faces=1)
    face_landmarker = vision.FaceLandmarker.create_from_options(options)
    print("✅ MediaPipe Face Landmarker initialized successfully!")
except Exception as e:
    print(f"❌ Error initializing MediaPipe: {e}")
    face_landmarker = None

# Initialize MediaPipe drawing utilities
mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles

# Load the ultra-optimized model and scaler
print("Loading ultra-optimized model...")
try:
    face_shape_model = joblib.load('Ultra_Optimized_RandomForest.joblib')
    scaler = joblib.load('ultra_optimized_scaler.joblib')
    print("✅ Ultra-optimized model loaded successfully!")
except Exception as e:
    print(f"❌ Error loading model files: {e}")
    face_shape_model = None
    scaler = None

def distance_3d(p1, p2):
    """Calculate 3D Euclidean distance between two points."""
    return np.linalg.norm(np.array(p1) - np.array(p2))

def smart_preprocess_image(image):
    """
    Smart image preprocessing to get the best face region.
    This addresses the issue of users not providing perfect images.
    """
    h, w = image.shape[:2]
    
    # First, try to detect face and get bounding box
    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb_image)
    detection_result = face_landmarker.detect(mp_image)
    
    if detection_result.face_landmarks:
        # Get face landmarks
        face_landmarks = detection_result.face_landmarks[0]
        
        # Calculate face bounding box with padding
        x_coords = [landmark.x for landmark in face_landmarks]
        y_coords = [landmark.y for landmark in face_landmarks]
        
        # Convert normalized coordinates to pixel coordinates
        x_min = int(min(x_coords) * w)
        x_max = int(max(x_coords) * w)
        y_min = int(min(y_coords) * h)
        y_max = int(max(y_coords) * h)
        
        # Add generous padding around face (40% padding for better context)
        face_width = x_max - x_min
        face_height = y_max - y_min
        pad_x = int(face_width * 0.4)  # 40% padding
        pad_y = int(face_height * 0.4)
        
        # Calculate crop coordinates
        x1 = max(0, x_min - pad_x)
        x2 = min(w, x_max + pad_x)
        y1 = max(0, y_min - pad_y)
        y2 = min(h, y_max + pad_y)
        
        # Crop the face region
        face_crop = image[y1:y2, x1:x2]
        
        # Resize to standard size while maintaining aspect ratio
        target_size = 224
        crop_h, crop_w = face_crop.shape[:2]
        
        # Calculate scale to fit in target size
        scale = min(target_size / crop_w, target_size / crop_h)
        new_w = int(crop_w * scale)
        new_h = int(crop_h * scale)
        
        # Resize maintaining aspect ratio
        resized = cv2.resize(face_crop, (new_w, new_h))
        
        # Create final image with padding to exact target size
        final_image = np.zeros((target_size, target_size, 3), dtype=np.uint8)
        
        # Center the resized image
        start_y = (target_size - new_h) // 2
        start_x = (target_size - new_w) // 2
        final_image[start_y:start_y + new_h, start_x:start_x + new_w] = resized
        
        return final_image
    else:
        # If no face detected, just resize to standard size
        return cv2.resize(image, (224, 224))

def extract_optimized_features(coords):
    """
    Extract optimized features for face shape detection.
    Uses only the most important landmarks for efficiency.
    """
    # Key landmarks for face shape analysis
    landmark_indices = {
        'forehead_top': 10,
        'forehead_left': 21,
        'forehead_right': 251,
        'cheek_left': 234,
        'cheek_right': 454,
        'jaw_left': 172,
        'jaw_right': 397,
        'chin': 152,
    }

    # Extract chosen points
    lm = {name: coords[idx] for name, idx in landmark_indices.items()}

    # Calculate key measurements
    face_height = distance_3d(lm['forehead_top'], lm['chin'])
    face_width = distance_3d(lm['cheek_left'], lm['cheek_right'])
    jaw_width = distance_3d(lm['jaw_left'], lm['jaw_right'])
    forehead_width = distance_3d(lm['forehead_left'], lm['forehead_right'])

    # Calculate ratios (scale-invariant features)
    width_to_height = face_width / face_height
    jaw_to_forehead = jaw_width / forehead_width
    jaw_to_face = jaw_width / face_width
    forehead_to_face = forehead_width / face_width

    # Additional shape features
    face_area = face_width * face_height
    jaw_angle = np.arctan2(lm['jaw_right'][1] - lm['jaw_left'][1], 
                          lm['jaw_right'][0] - lm['jaw_left'][0])
    
    # Return optimized feature vector
    features = np.array([
        width_to_height,
        jaw_to_forehead,
        jaw_to_face,
        forehead_to_face,
        face_area,
        jaw_angle
    ])

    return features

def get_face_shape_label(label):
    shapes = ["Heart", "Oval", "Round", "Square", "Oblong"]
    return shapes[label]

def draw_landmarks_on_image(rgb_image, detection_result):
    face_landmarks_list = detection_result.face_landmarks
    annotated_image = np.copy(rgb_image)

    for idx in range(len(face_landmarks_list)):
        face_landmarks = face_landmarks_list[idx]

        # Create landmark proto
        face_landmarks_proto = landmark_pb2.NormalizedLandmarkList()
        face_landmarks_proto.landmark.extend([
            landmark_pb2.NormalizedLandmark(x=landmark.x, y=landmark.y, z=landmark.z) for landmark in face_landmarks
        ])

        # Draw face landmarks
        mp_drawing.draw_landmarks(
            image=annotated_image,
            landmark_list=face_landmarks_proto,
            connections=mp.solutions.face_mesh.FACEMESH_TESSELATION,
            landmark_drawing_spec=None,
            connection_drawing_spec=mp_drawing_styles.get_default_face_mesh_tesselation_style())
        mp_drawing.draw_landmarks(
            image=annotated_image,
            landmark_list=face_landmarks_proto,
            connections=mp.solutions.face_mesh.FACEMESH_CONTOURS,
            landmark_drawing_spec=None,
            connection_drawing_spec=mp_drawing_styles.get_default_face_mesh_contours_style())
        mp_drawing.draw_landmarks(
            image=annotated_image,
            landmark_list=face_landmarks_proto,
            connections=mp.solutions.face_mesh.FACEMESH_IRISES,
            landmark_drawing_spec=None,
            connection_drawing_spec=mp_drawing_styles.get_default_face_mesh_iris_connections_style())

    return annotated_image

def allowed_file(filename):
        return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']

def generate_frames():
        cap = cv2.VideoCapture(0)
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb_frame)
            detection_result = face_landmarker.detect(image)

            if detection_result.face_landmarks:
                for face_landmarks in detection_result.face_landmarks:
                    landmarks = [[lm.x, lm.y, lm.z] for lm in face_landmarks]
                    landmarks = np.array(landmarks)
                    face_features = extract_optimized_features(landmarks)
                    
                    # Normalize features using the scaler
                    face_features_scaled = scaler.transform(face_features.reshape(1, -1))
                    
                    face_shape_label = face_shape_model.predict(face_features_scaled)[0]
                    face_shape = get_face_shape_label(face_shape_label)
                    annotated_image = draw_landmarks_on_image(rgb_frame, detection_result)
                    cv2.putText(annotated_image, f"Face Shape: {face_shape}", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            else:
                annotated_image = rgb_frame

            ret, buffer = cv2.imencode('.jpg', annotated_image)
            frame = buffer.tobytes()
            yield (b'--frame\r\n'
                b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/health')
def health_check():
    """Health check endpoint for Hugging Face Spaces"""
    status = {
        "status": "healthy",
        "mediapipe_loaded": face_landmarker is not None,
        "ml_model_loaded": face_shape_model is not None and scaler is not None,
        "cloudinary_available": os.getenv("CLOUDINARY_CLOUD_NAME") is not None,
        "mongodb_available": os.getenv("MONGO_URI") is not None
    }
    return jsonify(status)

@app.route('/analyze', methods=['POST'])
def analyze_face():
    # Check if models are loaded
    if face_landmarker is None:
        return jsonify({"error": "MediaPipe model not loaded"}), 500
    if face_shape_model is None or scaler is None:
        return jsonify({"error": "ML model not loaded"}), 500
    
    if 'file' not in request.files:
        return jsonify({"error": "No file part"}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    try:
        # --- 1. Read image and smart preprocessing ---
        img_bytes = file.read()
        nparr = np.frombuffer(img_bytes, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        # Smart preprocessing: detect face and crop optimally
        processed_img = smart_preprocess_image(img)
        
        # Convert to RGB for MediaPipe
        rgb_image = cv2.cvtColor(processed_img, cv2.COLOR_BGR2RGB)
        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb_image)
        
        detection_result = face_landmarker.detect(mp_image)

        if not detection_result.face_landmarks:
            return jsonify({"error": "No face detected"}), 400

        # --- 2. Get data, calculate features, and predict shape ---
        face_landmarks = detection_result.face_landmarks[0]
        
        # First, calculate the optimized features
        landmarks_normalized = np.array([[lm.x, lm.y, lm.z] for lm in face_landmarks])
        face_features = extract_optimized_features(landmarks_normalized)
        
        # Normalize features using the scaler
        face_features_scaled = scaler.transform(face_features.reshape(1, -1))
        
        # Then, predict the shape using calibrated features
        face_shape_label = face_shape_model.predict(face_features_scaled)[0]
        face_shape = get_face_shape_label(face_shape_label)
        
        # Get confidence scores
        confidence_scores = face_shape_model.predict_proba(face_features_scaled)[0]
        confidence = confidence_scores[face_shape_label]

        # --- 3. Draw landmarks on the image ---
        annotated_image_rgb = draw_landmarks_on_image(rgb_image, detection_result)
        # cv2.putText(annotated_image_rgb, f"Face Shape: {face_shape}", (20, 50), 
        #            cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        # cv2.putText(annotated_image_rgb, f"Confidence: {confidence:.3f}", (20, 90), 
        #            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

        # --- 4. Upload the PROCESSED image to Cloudinary ---
        annotated_image_bgr = cv2.cvtColor(annotated_image_rgb, cv2.COLOR_RGB2BGR)
        _, buffer = cv2.imencode('.jpg', annotated_image_bgr)
        processed_image_url = upload_image_to_cloudinary(buffer.tobytes())
        
        if not processed_image_url:
            return jsonify({"error": "Failed to upload processed image"}), 500

        # --- 5. Calculate Measurements using CALIBRATED values ---
        landmarks_normalized = np.array([[lm.x, lm.y, lm.z] for lm in face_landmarks])

        # Define more accurate landmark points for measurements
        p_iris_l = landmarks_normalized[473] # Left Iris
        p_iris_r = landmarks_normalized[468] # Right Iris
        
        p_forehead_top = landmarks_normalized[10]  # Top of forehead hairline
        p_chin_tip = landmarks_normalized[152]     # Bottom of chin

        p_cheek_l = landmarks_normalized[234]      # Left cheekbone edge
        p_cheek_r = landmarks_normalized[454]      # Right cheekbone edge

        p_jaw_l = landmarks_normalized[172]        # Left jaw point
        p_jaw_r = landmarks_normalized[397]        # Right jaw point

        p_forehead_l = landmarks_normalized[63]   # Left forehead edge
        p_forehead_r = landmarks_normalized[293]  # Right forehead edge
        
        # IPD-based calibration
        AVG_IPD_CM = 6.3
        dist_iris = distance_3d(p_iris_l, p_iris_r)
        cm_per_unit = AVG_IPD_CM / dist_iris if dist_iris != 0 else 0

        # Calculate all distances
        dist_face_length = distance_3d(p_forehead_top, p_chin_tip)
        dist_cheek_width = distance_3d(p_cheek_l, p_cheek_r)
        dist_jaw_width = distance_3d(p_jaw_l, p_jaw_r)
        dist_forehead_width = distance_3d(p_forehead_l, p_forehead_r)

        # Convert to cm and apply calibration adjustments
        face_length_cm = (dist_face_length * cm_per_unit) + 5.0      # +4cm calibration (increased from +2cm)
        cheekbone_width_cm = (dist_cheek_width * cm_per_unit) + 4.0  # +3cm calibration (increased from +2cm)
        jaw_width_cm = (dist_jaw_width * cm_per_unit) +0.5                 # No calibration (already accurate)
        forehead_width_cm = (dist_forehead_width * cm_per_unit) + 6.0 # +5cm calibration (increased from +3.5cm)
        
        # Jaw curve ratio is a relative measure, so it doesn't need cm conversion
        jaw_curve_ratio = dist_face_length / dist_cheek_width if dist_cheek_width != 0 else 0

        measurements = {
            "face_length_cm": float(face_length_cm),
            "cheekbone_width_cm": float(cheekbone_width_cm),
            "jaw_width_cm": float(jaw_width_cm),
            "forehead_width_cm": float(forehead_width_cm),
            "jaw_curve_ratio": float(jaw_curve_ratio)
        }

        # --- 6. Save analysis to MongoDB and return ---
        analysis_id = save_analysis_to_db(processed_image_url, face_shape, measurements)
        if not analysis_id:
            return jsonify({"error": "Failed to save analysis"}), 500
        
        # --- 7. Return the complete result ---
        return jsonify({
            "message": "Analysis successful",
            "analysis_id": analysis_id,
            "image_url": processed_image_url,
            "face_shape": face_shape,
            "confidence": float(confidence),
            "all_probabilities": {
                "Heart": float(confidence_scores[0]),
                "Oval": float(confidence_scores[1]),
                "Round": float(confidence_scores[2]),
                "Square": float(confidence_scores[3]),
                "Oblong": float(confidence_scores[4])
            },
            "measurements": measurements,
            "calibration_applied": {
                "face_length_adjustment": "+4.0cm (increased from +2.0cm)",
                "forehead_width_adjustment": "+5.0cm (increased from +3.5cm)", 
                "cheekbone_width_adjustment": "+3.0cm (increased from +2.0cm)",
                "jaw_width_adjustment": "none (already accurate)",
                "note": "Calibration adjustments increased based on user feedback"
            }
        })

    except Exception as e:
        print(f"An error occurred: {e}")
        return jsonify({"error": f"An error occurred: {str(e)}"}), 500

@app.route('/video_feed')
def video_feed():
    return Response(generate_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')

@app.route('/real_time')
def real_time():
    return render_template('real_time.html')

if __name__ == '__main__':
    # Get port from environment variable (required for Hugging Face Spaces)
    port = int(os.environ.get('PORT', 7860))
    # Run on 0.0.0.0 to accept connections from any IP (required for HF Spaces)
    app.run(host='0.0.0.0', port=port, debug=False)