Spaces:
Runtime error
Runtime error
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) |