from flask import Flask, request, jsonify
from huggingface_hub import from_pretrained_keras
import numpy as np
from PIL import Image
import io

app = Flask(__name__)

# Load the model
# model = from_pretrained_keras("MissingBreath/recycle-garbage-model")
model = from_pretrained_keras("./recycle-garbage-model")

# Class labels
# class_labels = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']

@app.route('/classify', methods=['POST'])
def classify():
    file = request.files['image']
    if file:
        img = Image.open(io.BytesIO(file.read()))
        img = img.resize((128, 128))
        img_array = np.array(img) / 255.0
        img_array = np.expand_dims(img_array, axis=0)
        predictions = model.predict(img_array)
        predicted_class_idx = np.argmax(predictions)
        # predicted_class = class_labels[predicted_class_idx]
        # return jsonify({'prediction': predicted_class})
        return jsonify({'prediction': predicted_class_idx})
    else:
        return jsonify({'error': 'No image provided'}), 400

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