Update model.py

model.py

@@ -1,27 +1,32 @@
-import gradio as gr
+import tensorflow as tf
+import numpy as np
 from PIL import Image
-from model import predict
-
-def classify_image(img: Image.Image):
-    label, confidence, probs = predict(img)
-
-    return (
-        label,
-        round(confidence, 3),
-        {k: round(v, 3) for k, v in probs.items()}
-    )
-
-demo = gr.Interface(
-    fn=classify_image,
-    inputs=gr.Image(type="pil", label="Upload an image"),
-    outputs=[
-        gr.Label(label="Predicted Class"),
-        gr.Number(label="Confidence"),
-        gr.JSON(label="All Probabilities")
-    ],
-    title="Animal Image Classifier",
-    description="Upload an image and the model will predict the animal."
+import os
+
+MODEL_PATH = os.path.join(
+    os.path.dirname(__file__),
+    "saved_model",
+    "Inception_V3_Animals_Classification.h5"
 )
 
-if __name__ == "__main__":
-    demo.launch()
+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