from flask import Flask, request, jsonify
import cv2
import joblib
import mediapipe as mp
import numpy as np
import tempfile

app = Flask(__name__)

# Load model and label encoder
model = joblib.load("pose_classifier.joblib")
label_encoder = joblib.load("label_encoder.joblib")

# Initialize MediaPipe Pose
mp_pose = mp.solutions.pose
pose = mp_pose.Pose()

def predict_pose_from_image(image_bytes):
    # Convert image bytes to numpy array
    nparr = np.frombuffer(image_bytes, np.uint8)
    frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    if frame is None:
        return None, "Invalid image"

    # Convert to RGB
    img_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # Run MediaPipe Pose
    results = pose.process(img_rgb)

    if results.pose_landmarks:
        landmarks = results.pose_landmarks.landmark
        pose_data = [j.x for j in landmarks] + [j.y for j in landmarks] + \
                    [j.z for j in landmarks] + [j.visibility for j in landmarks]

        pose_data = np.array(pose_data).reshape(1, -1)
        y_pred = model.predict(pose_data)
        predicted_label = label_encoder.inverse_transform(y_pred)[0]
        return predicted_label, None
    else:
        return None, "No pose detected"

@app.route('/predict-pose', methods=['POST'])
def predict_pose():
    if 'frame' not in request.files:
        return jsonify({"error": "No image frame uploaded"}), 400

    file = request.files['frame']
    img_bytes = file.read()

    label, error = predict_pose_from_image(img_bytes)
    if error:
        return jsonify({"error": error}), 400

    return jsonify({"predicted_pose": label})

if __name__ == "__main__":
    app.run(debug=True, port=5007)

# curl -X POST -F "frame=@your_image.jpg" http://localhost:5007/predict-pose