app/model.py

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