from flask import Flask, request, jsonify
from fastai.vision.all import *
from huggingface_hub import from_pretrained_fastai
from PIL import Image
import io

app = Flask(__name__)

def classify_image(image_file):
    # Load the trained model from Huggingface
    learn = from_pretrained_fastai("devdatanalytics/commonbean")
    
    # Open the image file
    img = Image.open(image_file)
    
    # Perform any necessary preprocessing on the image
    # For example, resizing or normalization
    img = img.resize((224, 224))  # Resize the image to match the model's input size
    
    # Convert the image to a BytesIO object
    img_bytes = io.BytesIO()
    img.save(img_bytes, format='PNG')
    img_bytes.seek(0)
    
    # Perform the classification
    pred_class, pred_idx, probs = learn.predict(img_bytes)
    
    # Return the classification results
    return f"Predicted Class: {pred_class}, Probability: {probs[pred_idx]:.2f}"

@app.route('/classify', methods=['POST'])
def classify():
    if 'image' not in request.files:
        return "No image file found", 400
    
    image_file = request.files['image']
    
    # Perform image classification
    classification_results = classify_image(image_file)
    
    return jsonify(classification_results)

if __name__ == '__main__':
    app.run()