File size: 1,443 Bytes
630d6ba | 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 | import tensorflow as tf
import numpy as np
from PIL import Image
# Load your trained CNN model
model = tf.keras.models.load_model("saved_model/Sports_Balls_Classification.h5")
# Same label order you used when training (from LabelEncoder)
CLASS_NAMES = [
"american_football", "baseball", "basketball", "billiard_ball",
"bowling_ball", "cricket_ball", "football", "golf_ball",
"hockey_ball", "hockey_puck", "rugby_ball", "shuttlecock",
"table_tennis_ball", "tennis_ball", "volleyball"
]
def preprocess_image(img: Image.Image, target_size=(225, 225)):
"""
Preprocess a PIL image to match training pipeline:
- Convert to RGB
- Resize
- Convert to float32
- Normalize to [0,1]
- Add batch dimension
"""
img = img.convert("RGB") # ensure 3 channels
img = img.resize(target_size)
img = np.array(img).astype("float32") / 255.0 # normalize
img = np.expand_dims(img, axis=0) # (1, 225, 225, 3)
return img
def predict(img: Image.Image):
# Apply preprocessing
input_tensor = preprocess_image(img)
# Model prediction
preds = model.predict(input_tensor)
probs = preds[0]
class_idx = int(np.argmax(probs))
confidence = float(np.max(probs))
# Map all probabilities
prob_dict = {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}
return CLASS_NAMES[class_idx], confidence, prob_dict |