File size: 1,481 Bytes
af98bbd |
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 |
import numpy as np
from tensorflow.keras.applications.resnet50 import (
ResNet50,
decode_predictions,
preprocess_input,
)
from tensorflow.keras.preprocessing import image
# Load the model outside the function to ensure it's loaded only once
model = ResNet50(include_top=True, weights="imagenet")
def predict_image(img):
"""
Preprocesses an image and runs a pre-trained ResNet50 model to get a prediction.
Parameters
----------
img : PIL.Image
The image object to classify.
Returns
-------
class_name, pred_probability : tuple(str, float)
The model's predicted class as a string and the corresponding confidence
score as a number.
"""
# Resize the image to match model input dimensions (224, 224)
img = img.resize((224, 224))
# Convert Pillow image to np.array
x = image.img_to_array(img)
# Add an extra dimension for the batch size
x_batch = np.expand_dims(x, axis=0)
# Apply ResNet50-specific preprocessing
x_batch = preprocess_input(x_batch)
# Make predictions
predictions = model.predict(x_batch, verbose=0)
# Get predictions using model methods and decode predictions
top_pred = decode_predictions(predictions, top=1)[0][0] # imagenet_id, label, score
_, class_name, pred_probability = top_pred
# Convert probability to float and round it
pred_probability = round(float(pred_probability), 4)
return class_name, pred_probability
|