from flask import Flask, request, jsonify
from huggingface_hub import from_pretrained_fastai
from PIL import Image
import io
import os

app = Flask(__name__)

# Set the cache directory to a local path within the app
os.environ['HF_HOME'] = './.cache'

# Now import from huggingface_hub
from huggingface_hub import from_pretrained_fastai

learn = from_pretrained_fastai("hugginglearners/brain-tumor-detection-mri")

@app.route('/predict', methods=['POST'])
def predict():
    # Get image from the request
    file = request.files['file']
    img = Image.open(io.BytesIO(file.read()))
    
    # Make a prediction
    pred_class, _, probs = learn.predict(img)
    
    # Return the result as JSON
    return jsonify({
        'prediction': str(pred_class),
        'confidence': float(probs.max())
    })

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