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 = 'test1.png'  # Replace with the actual image path
result = classify_image(image_path)
print(result)