Update model.py

model.py

@@ -1,40 +1,32 @@
-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/Inception_V3_Animals_Classification.h5")
-
-# Same label order you used when training (from LabelEncoder)
-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")  # 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, 256, 256, 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))}
-
+import tensorflow as tf
+import numpy as np
+from PIL import Image
+import os
+
+MODEL_PATH = os.path.join(
+    os.path.dirname(__file__),
+    "saved_model",
+    "Inception_V3_Animals_Classification.h5"
+)
+
+model = tf.keras.models.load_model(MODEL_PATH)
+
+CLASS_NAMES = ["Cat", "Dog", "Snake"]
+
+def preprocess_image(img: Image.Image, target_size=(256, 256):
+    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)[0]
+
+    class_idx = int(np.argmax(preds))
+    confidence = float(np.max(preds))
+
+    prob_dict = {CLASS_NAMES[i]: float(preds[i]) for i in range(len(CLASS_NAMES))}
+
     return CLASS_NAMES[class_idx], confidence, prob_dict