Create inference.py

inference.py

@@ -0,0 +1,52 @@
+import tensorflow as tf
+import numpy as np
+from PIL import Image
+from typing import Dict
+
+# Load the ResNetV2 model
+model = tf.keras.models.load_model("resnetv2_model.h5")
+
+# Define the handler for the Inference API
+def predict(inputs: Dict) -> Dict:
+    """
+    Handle inference requests.
+    Args:
+        inputs (Dict): A dictionary with a key 'image' containing the base64-encoded image.
+    Returns:
+        Dict: A dictionary containing the predicted class label.
+    """
+    # Decode the image
+    if "image" not in inputs:
+        return {"error": "No image found in inputs"}
+    
+    # Preprocess the input image
+    image = preprocess_image(inputs["image"])
+    
+    # Perform inference
+    prediction = model.predict(image)
+    predicted_class = np.argmax(prediction, axis=1)[0]  # Get the predicted class index
+    
+    # Return the predicted class
+    return {"label": int(predicted_class)}
+
+def preprocess_image(image_base64: str) -> np.ndarray:
+    """
+    Preprocess the input image for ResNetV2.
+    Args:
+        image_base64 (str): Base64-encoded image.
+    Returns:
+        np.ndarray: Preprocessed image ready for inference.
+    """
+    from io import BytesIO
+    import base64
+
+    # Decode the base64 image
+    image_data = base64.b64decode(image_base64)
+    image = Image.open(BytesIO(image_data)).convert("RGB")
+    
+    # Resize and normalize the image
+    image = image.resize((224, 224))
+    image_array = np.array(image) / 255.0  # Normalize pixel values to [0, 1]
+    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
+    
+    return image_array