rohitashva's picture
Update app.py
d0f6bea verified
raw
history blame
1.77 kB
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