AhmedIbrahim007's picture
Create app.py
9c9c79e verified
raw
history blame
1.14 kB
from fastai.vision.all import *
# Load the pre-trained model
learn = load_learner('model.pkl')
# Define categories and map them to indices
searches = ['formal', 'casual', 'athletic']
searches = sorted(searches) # Ensure the categories are in sorted order
values = [i for i in range(0, len(searches))]
class_dict = dict(zip(searches, values))
def classify_image(image_path):
# Load the image from the provided path
img = PILImage.create(image_path)
# Make the prediction
classification, _, probs = learn.predict(img)
# Convert the prediction to a confidence dictionary
confidences = {label: float(probs[i]) for i, label in enumerate(class_dict)}
# If classification is not formal, return 'informal'
if classification != 'formal':
informal_confidence = sum(confidences[label] for label in class_dict if label != 'formal')
return {'informal': informal_confidence}
else:
return {'formal': confidences['formal']}
# Example usage with an image path
image_path = 'path_to_image.jpg' # Replace with the actual image path
result = classify_image(image_path)
print(result)