Spaces:
Sleeping
Sleeping
File size: 6,105 Bytes
db12d78 |
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 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
from keras.models import load_model
from PIL import Image
import numpy as np
import requests
from io import BytesIO
import timeit
import gdown
import json
import os
# Define the class labels for the model's output
# output_class = ["Plastic bottle/Can to deposit in Supermarkets", "Big Cardboard bin", "Unrecyclable garbage", "Glass - Purple bin", "Organic waste - Composter", "Grocery Packages - Orange bin", "Paper - Blue bin"]
with open('Bins.json', 'r') as f:
bins_dict = json.load(f)
output_class = bins_dict["output_class"]
url = "https://drive.google.com/uc?id=1TJSpf9wJYZSqIm5ytqXkU8RLuKJmmHHa"
output = 'vgg16_model_after_one_epoch.h5'
# check if path to file exist
if not os.path.exists(output):
gdown.download(url, output, quiet=False)
# Load the pre-trained model
model = load_model(output)
def predict(img_input, is_url=False):
"""
Function to predict the class labels and probabilities of an input image.
Args:
file (str): Path to the input image file or URL.
is_url (bool): Indicates whether the file argument is a URL or local file path.
Returns:
dict: A dictionary containing the predicted class labels as keys and
their corresponding probabilities as values.
"""
if isinstance(img_input, str):
if is_url is True:
# Download the image from the URL
response = requests.get(img_input)
img = Image.open(BytesIO(response.content))
else:
# Open the image file from the local file path
img = Image.open(img_input)
else:
img = img_input
# Resize the image to match the input size expected by the model (224x224)
resized_image = img.resize((224, 224))
# Normalize the pixel values to be between 0 and 1
normalized_image = np.array(resized_image) / 255.0
# Expand dimensions to match the input shape expected by the model
pred_image = np.expand_dims(normalized_image, axis=0)
# Make predictions using the model
predicted_array = model.predict(pred_image)
# Create a dictionary of predicted class labels and probabilities
pred_labels_and_probs = {
output_class[i]: float(predicted_array[0][i])
for i in range(len(output_class))
}
return pred_labels_and_probs
def image_identity(image):
"""
Take an image as input and return prediction information and an image of the predicted waste bin.
Parameters:
image (np.array): An image in numpy array format.
Returns:
bin_image1 (np.array): An image of the predicted waste bin.
predicted_class1 (str): The class of the predicted bin.
predicted_prob1 (float): The probability of the predicted class.
predicted_class2 (str): The second most likely class.
predicted_prob2 (float): The probability of the second most likely class.
prediction_time (str): The time taken to make the prediction.
"""
image = Image.fromarray((image * 255).astype(np.uint8))
# Load the classification and image data from the JSON file
with open('Bins.json', 'r') as f:
bins_dict = json.load(f)
# Retrieve the class labels and bin images from the dictionary
output_class = bins_dict["output_class"]
bin_images = bins_dict["bin_images"]
# Record the start time before making predictions
start_time = timeit.default_timer()
# Make predictions on the provided image
predicted_class1, predicted_prob1, predicted_class2, predicted_prob2 = check_item(image)
# Record the end time after making predictions and calculate the total time taken
end_time = timeit.default_timer()
pred_time = round(end_time - start_time, 2)
prediction_time = f"{pred_time} seconds"
# Retrieve the image URL of the predicted bin
bin_1 = output_class.index(predicted_class1)
bin_image_url1 = bin_images[bin_1]
response1 = requests.get(bin_image_url1)
# Open the image and convert to a numpy array for returning
bin_image1 = Image.open(BytesIO(response1.content))
bin_image1 = np.array(bin_image1)
return bin_image1, predicted_class1, predicted_prob1, predicted_class2, predicted_prob2, prediction_time
def check_item(file_path, is_url=False):
"""
Function to make predictions on an input image and compare the results with the expected classes.
Args:
file_path (str): Path to the input image file or URL.
is_url (bool): Indicates whether the file argument is a URL or local file path.
"""
# Obtain the predicted class labels and probabilities
prediction_dict = predict(file_path, is_url)
# Get the index of the highest prediction
ind = np.argmax(list(prediction_dict.values()))
predicted_class = list(prediction_dict.keys())[ind]
predicted_prob = list(prediction_dict.values())[ind]
predicted_prob = f"{predicted_prob:.2f}"
print(f"1st Prediction probabilities: {predicted_class} {predicted_prob}")
# Get the index of the second highest prediction
sorted_probs = np.sort(list(prediction_dict.values()))
ind2 = np.argsort(list(prediction_dict.values()))[-2]
predicted_class2 = list(prediction_dict.keys())[ind2]
predicted_prob2 = sorted_probs[-2]
predicted_prob2 = f"{predicted_prob2:.2f}"
print(f"2nd Prediction probabilities: {predicted_class2} {predicted_prob2}")
return predicted_class, predicted_prob, predicted_class2, predicted_prob2
if __name__ == "__main__":
# Prompt the user to enter the file path to the input image
file_path = input("Enter the file path or URL to the input image: ")
# Prompt the user to indicate if the input is a URL or a local file path
is_url = input("Is the input a URL? (y/n): ").lower() == "y"
# Measure the time taken to make predictions
start_time = timeit.default_timer()
# Make predictions and compare with expected results
check_item(file_path, is_url)
# Calculate and print the prediction time
end_time = timeit.default_timer()
pred_time = round(end_time - start_time, 2)
print("Prediction time:", pred_time, "seconds")
|