|
|
import tensorflow as tf
|
|
|
import numpy as np
|
|
|
from PIL import Image
|
|
|
|
|
|
|
|
|
model = tf.keras.models.load_model("saved_model/Inception_V3_Animals_Classification.h5")
|
|
|
|
|
|
|
|
|
CLASS_NAMES = ["Cat", "Dog", "Snake"]
|
|
|
|
|
|
def preprocess_image(img: Image.Image, target_size=(256, 256)):
|
|
|
"""
|
|
|
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")
|
|
|
img = img.resize(target_size)
|
|
|
img = np.array(img).astype("float32") / 255.0
|
|
|
img = np.expand_dims(img, axis=0)
|
|
|
return img
|
|
|
|
|
|
def predict(img: Image.Image):
|
|
|
|
|
|
input_tensor = preprocess_image(img)
|
|
|
|
|
|
|
|
|
preds = model.predict(input_tensor)
|
|
|
probs = preds[0]
|
|
|
|
|
|
class_idx = int(np.argmax(probs))
|
|
|
confidence = float(np.max(probs))
|
|
|
|
|
|
|
|
|
prob_dict = {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}
|
|
|
|
|
|
return CLASS_NAMES[class_idx], confidence, prob_dict |