File size: 4,559 Bytes
7c9e31a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | from typing import Dict, List, Any
import tensorflow as tf
import numpy as np
from PIL import Image
import io
import base64
import os
class EndpointHandler():
def __init__(self, path=""):
"""
Initialize the handler by loading the model and class labels.
Args:
path: Path to the model directory (provided by Hugging Face)
"""
# Load the trained model
model_path = os.path.join(path, "food_cassifier.h5")
self.model = tf.keras.models.load_model(model_path)
# Load class names from classes.txt
classes_path = os.path.join(path, "classes.txt")
with open(classes_path, 'r') as f:
self.class_names = [line.strip() for line in f.readlines()]
print(f"Model loaded successfully with {len(self.class_names)} classes")
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
Process the inference request.
Args:
data: Dictionary containing the input data
Expected format: {"inputs": "<base64_encoded_image>"}
or {"inputs": {"image": "<base64_encoded_image>"}}
Returns:
List of predictions with class labels and confidence scores
"""
# Extract the input image
inputs = data.get("inputs", "")
# Handle different input formats
if isinstance(inputs, dict):
image_data = inputs.get("image", "")
else:
image_data = inputs
# Decode base64 image
try:
# Remove data URL prefix if present
if isinstance(image_data, str) and "base64," in image_data:
image_data = image_data.split("base64,")[1]
# Decode base64
image_bytes = base64.b64decode(image_data)
image = Image.open(io.BytesIO(image_bytes))
except Exception as e:
return [{"error": f"Failed to decode image: {str(e)}"}]
# Preprocess the image
try:
preprocessed_image = self.preprocess_image(image)
except Exception as e:
return [{"error": f"Failed to preprocess image: {str(e)}"}]
# Make prediction
try:
prediction = self.model.predict(preprocessed_image)
except Exception as e:
return [{"error": f"Failed to make prediction: {str(e)}"}]
# Postprocess the prediction
results = self.postprocess_prediction(prediction)
return results
def preprocess_image(self, image: Image.Image) -> np.ndarray:
"""
Preprocess the input image for the model.
Args:
image: PIL Image object
Returns:
Preprocessed numpy array of shape (1, 224, 224, 3)
"""
# Convert to RGB if needed (handles RGBA, grayscale, etc.)
if image.mode != 'RGB':
image = image.convert('RGB')
# Convert to numpy array
img_array = tf.keras.preprocessing.image.img_to_array(image)
# Resize to (224, 224)
img_array = tf.image.resize(img_array, (224, 224))
# Expand dimensions to create batch: (224, 224, 3) -> (1, 224, 224, 3)
img_array = np.expand_dims(img_array, axis=0)
# Preprocess using MobileNetV2 preprocessing
img_array = tf.keras.applications.mobilenet_v2.preprocess_input(img_array)
return img_array
def postprocess_prediction(self, prediction: np.ndarray) -> List[Dict[str, Any]]:
"""
Postprocess the model prediction to return human-readable results.
Args:
prediction: Model output array of shape (1, 101)
Returns:
List containing prediction results with top predictions
"""
# Get the predicted class index
predicted_idx = int(np.argmax(prediction[0]))
confidence = float(prediction[0][predicted_idx])
predicted_label = self.class_names[predicted_idx]
# Get top 5 predictions
top_5_indices = np.argsort(prediction[0])[-5:][::-1]
top_5_predictions = [
{
"label": self.class_names[int(idx)],
"score": float(prediction[0][idx])
}
for idx in top_5_indices
]
# Return the results
return [
{
"predicted_class": predicted_label,
"predicted_index": predicted_idx,
"confidence": confidence,
"top_5_predictions": top_5_predictions
}
]
|