Spaces:
Sleeping
Sleeping
File size: 2,388 Bytes
f20fe1f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import numpy as np
import argparse
import os
# Define class labels (ensure these match your training)
CNN_CLASS_LABELS = ['Side Chest', 'Front Double Biceps', 'Back Double Biceps', 'Front Lat Spread', 'Back Lat Spread']
MODEL_PATH = 'bodybuilding_pose_classifier_savedmodel.keras' # Corrected model path
def predict_pose_from_image(model, img_path):
"""
Loads an image, preprocesses it, and predicts the bodybuilding pose.
Args:
model: The loaded Keras model.
img_path (str): Path to the image file.
Returns:
tuple: (predicted_class_label, confidence_score) or (None, None) if error.
"""
try:
if not os.path.exists(img_path):
print(f"Error: Image path not found: {img_path}")
return None, None
# Load and preprocess the image
img = image.load_img(img_path, target_size=(150, 150))
img_array = image.img_to_array(img)
img_array = np.expand_dims(img_array, axis=0) / 255.0 # Normalize
# Make prediction
predictions = model.predict(img_array)
predicted_class_index = np.argmax(predictions, axis=1)[0]
confidence = float(np.max(predictions))
predicted_class_label = CNN_CLASS_LABELS[predicted_class_index]
return predicted_class_label, confidence
except Exception as e:
print(f"Error during prediction: {e}")
return None, None
def main():
parser = argparse.ArgumentParser(description="Classify a bodybuilding pose from an image.")
parser.add_argument("image_path", help="Path to the input image file.")
args = parser.parse_args()
# Load the Keras model
print(f"Loading model from: {MODEL_PATH}")
try:
model = load_model(MODEL_PATH)
# Optional: Print model summary to verify
# model.summary()
except Exception as e:
print(f"Error loading model: {e}")
return
print(f"Classifying image: {args.image_path}")
predicted_pose, confidence_score = predict_pose_from_image(model, args.image_path)
if predicted_pose and confidence_score is not None:
print(f"Predicted Pose: {predicted_pose}")
print(f"Confidence: {confidence_score:.2f}")
if __name__ == "__main__":
main() |